Title: TabEBM: A Tabular Data Augmentation Method with Distinct Class-Specific Energy-Based Models
Venue: NeurIPS 2024
Accepted: yes
Keywords: tabular data

data augmentation

synthetic data generation

energy based model
Abstract: Data collection is often difficult in critical fields such as medicine, physics, and chemistry, yielding typically only small tabular datasets. However, classification methods tend to struggle with these small datasets, leading to poor predictive performance. Increasing the training set with additional synthetic data, similar to data augmentation in images, is commonly believed to improve downstream tabular 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 tabular 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 leads to improved classification performance across diverse datasets of various sizes, especially small ones. Code is available at https://github.com/andreimargeloiu/TabEBM.

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, yielding typically only small tabular datasets. However, classification methods tend to struggle with these small datasets, leading to poor predictive performance. Increasing the training set with additional synthetic data, similar to data augmentation in images, is commonly believed to improve downstream tabular classification performance. However, current tabular generative methods that learn either the joint distribution p ( x , y ) or the class-conditional distribution p ( x | 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 tabular 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 leads to improved classification performance across diverse datasets of various sizes, especially small ones. Code is available at https://github.com/andreimargeloiu/TabEBM .

1 Introduction

Many scientific domains within medicine, physics, and chemistry often rely on intricate and challenging data acquisition procedures [5, 50, 4, 32, 76, 12] that typically render smallsize tabular datasets [5, 46]. Using these to train machine learning models that can aid in tasks such as disease diagnosis [52, 37], material property prediction [35], and chemical compound classification [11], can lead to poor performance [74, 52, 37]. In the case of learning tasks which leverage image and text data, a standard remedy to address performance issues due to data scarcity is employing data augmentation techniques [72, 73, 60, 71] that generate additional synthetic samples from existing data.

-30

-20

-10

0

10

2000

1500

1000

500

0.2

0.3

0.4

0.5

1.0

0.9

0.8

0.7

0.6

1.0

0.8

0.6

0.4

0.2

1.0

0.8

0.6

0.4

0.2

InverseKL

Efficiency

TabEBM (Ours)

TabPFGen

ARF

GOGGLE

TabDDPM

Utility

NFLOW

CTGAN

TVAE

SMOTE

DCR

d-presence

EBM

EBM

Synthetic data

Real data

(input space)

TabEBM

Class-speci

fi

c EBMs

Class-speci

fi

c sampling

However, applying data augmentation to tabular data introduces additional challenges, as tabular datasets are often very diverse and lack explicit symmetries [8], such as rotations or translations seen in images. Consequently, existing tabular data augmentation methods often yield mixed results and can even degrade model performance [51, 71, 48], hindering their widespread adoption.

Tabular augmentation typically involves training generative models to approximate either the joint distribution p ( x , y ) [85, 24] or the class-conditional distribution p ( x | y ) [85, 42, 83, 47, 48]. A key challenge of joint distribution methods is maintaining the original training label distribution, as sampling from such generators can produce label distributions that deviate from the original and even fail to generate data for specific classes (see Appendix C for an example). These issues compromise the effectiveness of data augmentation [51] by undermining the label accuracy and distribution. On the other hand, while class-conditional models that learn p ( x | y ) preserve the stratification of the original data, they often employ a shared model to represent all class-conditional densities. This, however, can lead to overfitting, particularly in imbalanced datasets where the model may prioritise more frequent classes [21], ignoring unique features needed for generating label-invariant samples. Additionally, in datasets with limited data, this can lead to mode collapse [68, 70], where the model does not effectively capture the diversity of each class [70], and thus tends to perform poorly in a multi-class setting.

To address the challenges of class-conditional tabular generation, we introduce TabEBM (Figure 2), a new method for tabular data augmentation utilising Energy-Based Models (EBMs). Our method introduces two innovations: (i) Distinct class-specific models: TabEBM constructs a collection of individual models - one for each class - which, by design, enables learning distinct marginal distributions for the inputs associated with each class. This, in turn, enables performing data augmentation while maintaining the original label distribution. (ii) Generative models: we build novel class-specific generators that produce high-quality synthetic data even from extremely few samples. Specifically, we create a surrogate binary classification task for each class and fit it with a pre-trained tabular in-context classifier. We then convert the binary classifier into an EBM, a generative model, without additional training. Using class-specific EBMs makes the energy landscape more robust to class overlaps, compared to using a single shared EBM to approximate the class-conditional distribution.

Our contributions can be summarised as:

Technical: We propose TabEBM, which is the first generative method to create class-specific EBMs, learning the marginal distribution for each class separately.

Empirical: We present the first comprehensive analysis of tabular data augmentation across different dataset sizes and use cases beyond predictive performance. Our analysis compares TabEBM with eight leading tabular generative models across various datasets, demonstrating that TabEBM consistently improves data augmentation performance on small datasets, while our generated data demonstrates better statistical fidelity and privacy-preserving properties (Figure 1).

Library: We release TabEBM as an open-source library, available at https://github.com/ andreimargeloiu/TabEBM . Our library enables off-the-shelf data generation and data augmentation on any tabular dataset without requiring training. Further details are available in Appendix B.5.

2 TabEBM

Notation. We address classification problems with C classes, denoted by Y = { 1 , 2 , . . . , C } . Let { ( x ( i ) , y i ) } N i =1 represent a dataset of N samples, each being a D -dimensional vector x ( i ) ∈ R D , with a corresponding label y i ∈ Y . For each class c ∈ Y , we define X c = { x ( i ) | y i = c } as the subset of samples labelled with class c . Let f θ ( · ) denote a classifier. The expression f θ ( x )[ y ] represents the (unnormalised) logit assigned to the class y for the input x .

2.1 Preliminaries on Energy-Based Models

An Energy-Based Model (EBM) [43] defines a probability density function p θ ( x ) through an energy function E ( x ) . Specifically, the model posits that p ( x ) ∝ e -E ( x ) , where E ( x ) represents the unnormalised negative log-density of the input x . In this framework, lower energy values correspond to higher probability densities. This relationship allows EBMs to model distributions by learning to assign lower energy to more probable configurations of x and higher energy to less probable ones.

An important observation is that energy-based models can utilise the same model architectures as standard classification models [29]. Typically, the logits f θ ( x )[ y ] from a classification model define a discriminative distribution through the softmax function, expressed as p θ ( y | x ) = softmax( f θ ( x )[ y ]) . Intriguingly, these same logits can be reinterpreted to define an energy-based model for the joint distribution p ( x , y ) . This is achieved by setting the energy function to E ( x , y ) = -f θ ( x ) . Furthermore, the energy function for the marginal distribution p ( x ) is obtained by marginalising over p ( x , y ) , resulting in E ( x ) = -LogSumExp y ′ f θ ( x )[ y ′ ] .

Such an energy-based model, trained with EBM-specific protocols on multiple classes, is typically used as a classifier, as demonstrated on several computer vision tasks in [29]. In contrast, in this work our focus is the opposite: we propose employing trained classifiers, one for each specific class, as a generative energy-based model for the class-conditional distributions p ( x | y ) . We apply our TabEBM method for generative tasks on tabular data.

2.2 Distinct Class-Specific Energy-Based Models

TabEBM is a class-conditional generative model p ( x | y ) implemented using a set of EBMs, { E 1 ( x ) , E 2 ( x ) , . . . , E C ( x ) } . Our approach assumes that the class-conditional density p ( x | y = c ) is best modelled using its class-specific data X c . Thus, for each class c , we construct a class-specific EBM, E c ( x ) , using only the data from that class, X c , such that p ( x | y = c ) ∝ exp( -E c ( x )) .

We derive each class-specific EBM E c ( x ) by training a classifier on a novel task and reinterpreting its logits. Specifically, for each class c , we propose a surrogate binary classification task to determine if a sample belongs to class c by comparing X c against a set of surrogate negative samples X neg c , which we show in Figure 3. Specifically, we generate the negative samples at the corners of a hypercube in R D . For each dimension d , the coordinates of a negative sample are either α neg dist σ d or -α neg dist σ d , where α neg dist is a fixed constant and σ d is the standard deviation of dimension d . For example, in R 3 , a negative sample might have coordinates [ α neg dist σ 1 , α neg dist σ 2 , -α neg dist σ 3 ] . Placing the negative samples at the corners of a hypercube ensures they are easily distinguishable from the real data, which is crucial for an accurate energy function (see Appendix D.1.1). This placement is also robust to variations in the number and distance of the negative samples (see Appendices D.1.2 and D.1.3).

We create the combined dataset D c for the surrogate binary classification task by labelling X c as 1 and X neg c as 0:

real samples

c

negative samples

c

We then train a binary classifier f c θ ( · ) on D c and use it to construct the class-specific energy E c ( x ) for class c . To do this, we reinterpret the logits { f c θ ( x )[0] , f c θ ( x )[1] } of the trained binary classifier as components of an approximated joint distribution for the surrogate binary task:

Next, we derive the approximated distribution p c ( x ) by marginalisation:

For the binary classifier f c θ ( · ) in the surrogate binary classification, we use TabPFN [33], a pre-trained tabular in-context model. Note that TabPFN is intended for inference only, with no updates to its parameters (see Section 4 for more details about TabPFN). In this context, 'training' the TabPFN classifier is analogous to the K-Nearest Neighbour algorithm, which simply performs inference based on a training dataset provided to the model. We apply TabPFN multiple times on separate datasets {D 1 , D 2 , . . . , D C } to obtain multiple classifiers { f 1 θ , f 2 θ , . . . , f C θ } . In Section 3.4, we explore why reinterpreting TabPFN's logits, trained on our surrogate binary tasks, can be useful for estimating an energy function. We emphasise that TabEBM is a general method, capable of using any gradient-based classifier that computes logits (using Equation (3)), and is not limited to TabPFN.

Generating data with TabEBM involves two steps. First, we sample a class c from the empirical distribution c ∼ p ( y ) . Then, we sample a data point x from the conditional distribution x ∼ p ( x | y = c ) approximated by the class-specific energy-based model E c ( x ) , as outlined in Algorithm 1. We employ Stochastic Gradient Langevin Dynamics (SGLD) [84] to perform this sampling. SGLD is an efficient method for high-dimensional data, combining stochastic gradient descent (SGLD) with Langevin dynamics. The update rule for SGLD at each iteration is:

where a Gaussian noise term ϵ t introduces randomness into the sampling process, enhancing the exploration of the distribution. In practice, the step size and the noise standard deviations are often chosen separately, resulting in a biased sampler that allows for faster training. Appendix D.2 further shows that TabEBM is stable to hyperparameters for the sampling process.

In our method, SGLD performs iterative augmentation. We start by sampling close to real data and iteratively adjust these synthetic data points, steering them towards regions of higher probability under the learned energy model. TabEBM enables sampling from any specified class distribution, including the original class distribution, which is crucial for data augmentation.

Algorithm 1 TabEBM sampling from Class-Specific EBM E c ( x )

Input: Training data X c for class c , step size α step , noise scale α noise , initial perturbation σ start , number of steps T Output: Set of synthetic samples for class c

Initialise a surrogate binary classification task and train the model

1: Assign new labels to the samples X c from class c , setting them to class 1

3: Train a binary classifier f c θ on the dataset D c = ( X c ∪ X neg c , { 1 } |X c | ∪ { 0 } |X neg c | ) Synthesise samples using Stochastic Gradient Langevin Dynamics (SGLD)

2: Generate a set of surrogate negative samples X neg c and assign them class 0 labels

4: Initialise synthetic data points x synth 0 by sampling from N ( X c , σ 2 start I )

6: E c ( x synth t ) = -log ( exp( f c θ ( x synth t )[0]) + exp( f c θ ( x synth t )[1]) )

5: for each iteration t = 0 , 1 , . . . , T -1 do

7: x synth t +1 = x synth t -α step ∇ E c ( x synth t ) + N (0 , α 2 noise I )

9: return x synth T as the generated synthetic data for class c

8: end for

3 Experiments

We evaluate TabEBM by focusing on four research questions:

Data Augmentation Improvement (Q1, Section 3.1): Can TabEBM generate synthetic data that improves the accuracy of downstream predictors via data augmentation?

Privacy Preservation (Q3, Section 3.3): Can TabEBM generate synthetic data that finds a competitive trade-off between downstream performance and privacy preservation?

Statistical Fidelity (Q2, Section 3.2): Can TabEBM generate synthetic data with high statistical fidelity (i.e., with similar distributions to those of real data)?

Understanding TabEBM's energy formulation (Q4, Section 3.4): Why is TabEBM's classspecific energy effective, and how do the proposed surrogate tasks influence this?

Datasets. We utilise eight open-source tabular datasets from OpenML [7] across five domains: Medicine, Chemistry, Engineering, Language and Economics. As TabPFN utilises many small-size OpenML datasets in its meta-validation [33], it can lead to data leakage when evaluating TabEBM. Therefore, to provide fair comparisons, we select six additional leakage-free datasets from UCI [22]. These diverse datasets contain 7 to 77 features and 698 to 5500 samples across 2 to 26 classes. Five datasets contain both numerical and categorical features, while the remaining are numerical only. We further enlarge the evaluation scope by varying the degrees of data availability (i.e., N real ), leading up to 33 different test cases for the eight OpenML datasets. Appendix B.1 provides detailed descriptions.

Benchmark generators. Wecompare TabEBM against eight existing tabular data generation methods of eight different categories: (i) a standard interpolation method SMOTE [13]; (ii) a Variational Autoencoders (VAE) based method TVAE [85]; (iii) a Generative Adversarial Networks (GAN) method CTGAN [85]; (iv) a normalising flow model Neural Spine Flows (NFLOW) [24]; (v) a diffusion model TabDDPM [42]; (vi) a tree-based method Adversarial Random Forests (ARF) [83]; (vii) a Graph Neural Network (GNN) based method GOGGLE [47]; and (viii) a Prior-Data Fitted Networks (PFN) based method TabPFGen [48]. Furthermore, we also include a 'Baseline' model, where no data augmentation is applied (i.e., only real data is used to train downstream predictors). In Appendix B.6, we detail the settings used for TabEBM and all other generators.

Downstream predictors. We select six representative downstream predictors, including three standard baselines: Logistic Regression (LR) [16], KNN [27] and MLP [28]; two tree-based methods: Random Forest (RF) [10] and XGBoost [14]; and a PFN method: TabPFN [33].

General experimental setup. For each dataset of N samples, we first split it into stratified train and test sets. We create large test sets to reduce the likelihood that the model's performance is accidentally inflated due to a small, unrepresentative set of samples [69], and thus the test size is computed via N test = min ( N 2 , 500 ) . The full train set approximates the upper bound of the quality of synthetic data, and we call this set 'oracle'. We subsample the full train set to simulate different levels of data availability, thus the subset size N real varies over { 20 , 50 , 100 , 200 , 500 } . We split each subset into stratified training and validation sets with a ratio of 4:1. We provide detailed descriptions of data splitting in Appendix B.2 and preprocessing in Appendix B.3. We repeat the splitting ten times, summing up to 10 runs per subset size. The reported results are averaged by default over ten runs on the test sets. When aggregating results across datasets, we use the average distance to the minimum (ADTM) metric via affine renormalisation between the top-performing and worse-performing models [30, 54]. We provide the evaluation results averaged over six downstream predictors for a general conclusion, and the fine-grained numerical results for each predictor are in Appendix D.

Data augmentation setup. Given N real real samples, we first train generators on the real training data and then generate N syn synthetic samples. For training the downstream predictors, we expand the real training split by adding the synthetic samples. The real validation data is used for early stopping, and the real test set is used for evaluating the predictor's performance. The optimal N syn remains an open problem for tabular data [51, 71, 31]. Prior works [47, 48] mainly use synthetic sets with equivalent sizes to the real sets (i.e., N real = N syn ). However, we observe that N real = N syn can lead to highly unstable results, especially on small datasets that we investigate. Recent work has used different N syn for various generators, such as by applying post-processing [31, 71]. In this work, we want to provide a head-to-head comparison of the effect of data augmentation across subsampled datasets of varying sizes N real ∈ { 20 , 50 , 100 , 200 , 500 } . Therefore, we perform data augmentation with a large synthetic set ( N syn = 500 ) across all splits, and the synthetic data has the same class distribution as the real training data. We provide an illustrative figure of the data splitting setup in Appendix B.2.

3.1 Data Augmentation Improvement (Q1)

We evaluate the effect of using synthetic data for data augmentation by comparing the balanced accuracy of downstream predictors before and after augmentation. Typically, higher classification accuracy (i.e., ACCGenerator ) and accuracy improvements (i.e., ACCGenerator -ACCBaseline > 0 ) demonstrate the effectiveness of the synthetic data for data augmentation.

TabEBM effectively improves downstream performance across sample sizes, especially for very low-sample-size regimes. Table 1 and Figure 4 (Left) show that TabEBM exhibits competitive performance in data augmentation, generally achieving the highest downstream accuracy and average rank across most datasets and sample sizes. Notably, TabEBM is the only generator that consistently improves performance across sample sizes. A key observation is that most modern benchmark generators underperform even the Baseline, indicating poor approximated distributions in the low-sample-size regime. Moreover, TabEBM achieves the largest overall performance improvement on six leakage-free UCI datasets, further supporting its effectiveness (see Appendix D.5.2 for details).

Furthermore, TabEBM is the most widely applicable method among the top three competitive generators on the considered datasets: (i) SMOTE requires at least two samples per class for interpolation, and thus it is not applicable for some datasets, such as the 'protein' dataset ( N real = 20 ); (ii) TabPFGen cannot scale up to more than ten classes, such as the 'collins' dataset. In addition, TabEBM can stabilise downstream performance, especially when real data is very scarce ( N real = 20 ): TabEBM leads to smaller standard deviations than Baseline on seven out of eight datasets.

TabEBM effectively improves downstream performance across any number of classes, especially for more than ten classes. Figure 4 (Right) shows that TabEBM consistently outperforms the

0

100

200

300

400

500

Number of real samples

-

25

0

25

Normalised accuracy improvement (%)

Baseline

SMOTE

TVAE

CTGAN

NFLOW

TabDDPM

ARF

GOGGLE

TabPFGen

TabEBM (Ours)

At most 10

More than 10

Number of classes

-50

-25

0

25

-0.3

-1.0

-11.9

-29.0

-15.6

-28.5

-19.5

-30.1

-13.2

-10.5

-12.9

-28.8

-18.7

-17.3

+4.4

0.0

+5.9

+13.7

SMOTE

TVAE

CTGAN

NFLOW

TabDDPM

ARF

GOGGLE

TabPFGen

TabEBM (Ours)

Baseline with notable improvements, particularly in datasets with more than ten classes. In contrast, an increased number of classes tends to cause a performance degradation in the benchmark generators.

TabEBM is robust on imbalanced datasets. For the three binary OpenML datasets (i.e., 'biodeg', 'steel' and 'stock'), we adjust the class distribution in the training data to vary the class imbalance, while keeping the test data fixed. Figure 5 shows that TabEBM consistently outperforms Baseline, while the other generators exhibit performance degradation as data imbalance increases.

TabEBM is computationally efficient. Figure 6 shows the trade-off between accuracy and the time needed for generating stratified synthetic data (for data augmentation). We measure the total duration of (i) training the model and (ii) generating 500 synthetic samples. The results show that TabEBM is practical, as it achieves higher downstream accuracy with lower time costs.

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

-

5

.

0

-

2

.

5

0

.

0

2

.

5

5

.

0

7

.

5

10

.

0

Normalised accuracy improvement (%)

SMOTE

TVAE

CTGAN

TabDDPM

TabPFGen

TabEBM (Ours)

#Minority

:

(#Minority

+

#Majority)

Ratio

3.2 Statistical Fidelity (Q2)

We evaluate the fidelity of synthetic data by measuring the similarity of synthetic data to real train data and to real test data (Figure 7). We evaluate this similarity via (i) average inverse of the Kullback-Leibler Divergence (inverse KL) [17], (ii) p-value of Kolmogorov-Smirnov test (KS test) [39] and (iii) p-value of Chi-squared test ( χ 2 test) [55]. For full numerical results, including χ 2 test, see Appendix D.6. For all three metrics, a bigger value denotes that synthetic data is more likely to have the same distribution as real data.

In Figure 7 (a1&a2), TabEBM consistently exhibits the highest accuracy and distribution similarity between real train data and synthetic data, indicating that TabEBM learns the distributions of real train data better than benchmark generators. In Appendix D.6, we further show that TabEBM remains the most competitive method in similarity between real test data and synthetic data. This indicates that TabEBM can extrapolate beyond real train data and thus generate synthetic data from a more

0

500

1000

1500

2000

Data augmentation time (s)

-

20

-

10

0

10

Normalised accuracy improvement (%)

SMOTE

TVAE

CTGAN

NFLOW

TabDDPM

ARF

GOGGLE

TabPFGen

TabEBM (Ours)

↓

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

(a1)

Inverse KL (train, syn)

-

20

-

10

0

10

Normalised accuracy improvement (%)

SMOTE

TVAE

CTGAN

NFLOW

TabDDPM

ARF

GOGGLE

TabPFGen

TabEBM (Ours)

0

.

6

0

.

7

0

.

8

0

.

9

1

.

0

(a2)

KS test (train, syn)

-

20

-

10

0

10

SMOTE

TVAE

CTGAN

NFLOW

TabDDPM

ARF

GOGGLE

TabPFGen

TabEBM (Ours)

0

.

2

0

.

3

0

.

4

0

.

5

(b1)

DCR (train, syn)

-

40

-

20

0

10

Normalised accuracy change (%)

SMOTE

TVAE

CTGAN

NFLOW

TabDDPM

ARF

GOGGLE

TabPFGen

TabEBM (Ours)

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

(b2)

d

-presence (train, syn)

-

40

-

20

0

10

SMOTE

TVAE

CTGAN

NFLOW

TabDDPM

ARF

GOGGLE

TabPFGen

TabEBM (Ours)

↑

↑

↑

↓

general distribution that aligns with both train and test data. This extrapolation ability also explains why TabEBM can outperform Baseline via data augmentation (Section 3.1).

3.3 Privacy Preservation (Q3)

More broadly, data privacy is a critical concern for organisations and governments handling sensitive data [75]. Privacy-preserving synthetic data allows researchers and practitioners to bypass ethical and logistical issues while enabling model training and testing [38]. We further explore the use of TabEBM-generated data for data sharing, where only synthetic data is accessible for downstream users [75, 89, 86, 23, 42]. In this case, downstream models are trained exclusively on synthetic data.

Specifically, we evaluate synthetic data via three metrics: (i) balanced accuracy of downstream predictors trained with only synthetic data (i.e., train-on-synthetic, test-on-real [85, 42, 87]); (ii) median Distance to Closest Record (DCR) [88], where a greater DCR denotes synthetic data is less likely to be copied from real data; and (iii) δ -presense [62], where a smaller value denotes a lower re-identification risk for real data from synthetic data. Full numerical results are in Appendix D.7.

Figure 7 (b1&b2) shows that TabEBM consistently finds a better trade-off between accuracy and privacy preservation. Notably, the 'train-on-synthetic, test-on-real' scenario poses a greater challenge for generators in achieving high accuracy because real data is inaccessible for model training and data augmentation. Despite this difficulty, TabEBM is the only generator that surpasses the overall performance of training on real data (i.e., Baseline). The relatively high DCR for TabEBM indicates that it can extrapolate beyond real train data, aligning with the finding that TabEBM's synthetic data is statistically similar to real test data (Section 3.2). These results further suggest that TabEBM learns the general distribution of real data, and can generate high-quality synthetic data suitable for various purposes, including privacy preservation.

3.4 Why is TabEBM effective for estimating Energy-Based Models? (Q4)

Having established that TabEBM excels in data augmentation, we explore why classifier logits can be useful when reinterpreted as a class-conditional energy function. Figure 8 shows the logit distribution of TabPFN trained on surrogate binary tasks and the corresponding energy function of TabEBM (with TabPFN as the binary classifier) as the Euclidean distance from the real data increases.

We found it essential to place the negative samples far from the real data, since TabPFN, which is pre-trained to approximate Bayesian inference [33], has its confi-

0

5

10

15

20

Distance from the real data

0.0

2.5

5.0

7.5

10.0

12.5

15.0

Logit value

Logit

f

(

x

)[1] for real data

Logit

f

(

x

)[0] for negative samples

0

5

10

15

20

Distance from the real data

0

25

50

75

100

125

150

Unnorm. density (x1000)

TabEBM unnormalised density

dence influenced by the distance from the training data [53]. Figure 8 (left) shows that TabPFN outputs high logit values near the real data. As the distance from the real data increases, the logit f ( x )[1] decreases smoothly until the two logits become similar, making the classifier uncertain (because the class probabilities become equal). Figure 8 (right) shows that TabEBM's inferred density drops significantly as the maximum logit decreases, because p c ( x ) ∝ (exp( f ( x )[0]) + exp( f ( x )[1])) from Equation (3). Since SGLD sampling performs gradient ascent on the density, the TabEBM-generated samples will be close to the real data. These findings are consistent across datasets (see Appendix D.3), where TabPFN's logits remain positive, with similar ranges and a relatively constant sum as distance increases, warranting further investigation. Overall, TabPFN's distance-based uncertainty is useful for inferring accurate energy functions within our TabEBM framework. Since TabEBM can be paired with any other gradient-based classifier that produces logits, we leave these extensions for future work.

4 Discussion & Related Work

Section 3 showed that TabEBM efficiently generates high-fidelity data that can effectively improve the downstream performance via data augmentation. In Table 2, we further provide a summary of tabular data generative models analysed from three important perspectives: (i) Training: the type of distribution that the generators learn (crucial for preserving the original training label distribution), and the training costs associated with learning; (ii) Generation: do the generators employ class-specific models (reflecting their capability to capture unique features essential for label-invariant generation), and do models support stratified generation (crucial for effective data augmentation); (iii) Practicability: the scalability of the generators with respect to the number of classes (a common requirement in real-world multi-class tasks), and consistent downstream performance improvement across different class sizes.

Generative Models for Tabular Data. The common paradigm for tabular data generation is to adapt Generative Adversarial Networks (GANs) and Variational Autoencoders (VAEs) [85, 63]. For instance, TableGAN employs a convolutional neural network to optimise the label quality [63], and TVAEis introduced in [85] as a variant of V AE for tabular data. However, these methods learn the joint distribution and thus cannot preserve the stratification of the original data (Appendix C). CTGAN [85] refines the generation to be class-conditional. The recent ARF [83] is an adversarial variant of random forest for density estimation, and GOGGLE [47] enhances VAE by learning relational structure with a Graph Neural Network (GNN). Some recent work focuses on generation with denoising diffusion models [42, 87, 40, 44]. For instance, TabDDPM [42] demonstrates that diffusion models can approximate typical distributions of tabular data. Although these class-conditional models can preserve the label distribution, they struggle to outperform Baseline and standard SMOTE in data augmentation [71, 48].

We attribute the performance degradation in current class-conditional models to their reliance on a single shared model to approximate all class-conditional densities. For instance, another promising generative approach uses pre-trained models like Prior-Data Fitted Networks (PFNs), and the recent TabPFGen [48] adapts such models into one shared class-conditional generator. However, TabPFGen's shared generator can lead to inaccurate density estimates, particularly in high-noise and classimbalance situations (see examples in Appendix C). As noise increases, TabPFGen's inferred densities fluctuate significantly and diverge from the true data distributions. In contrast, TabEBM uses classspecific EBMs to model each class's marginal distributions, and the results in Appendix C reveal that our design choice reduces the impact of noise and data imbalance. TabEBM focuses on approximating and generating for one class at a time, remaining unaffected by noise from other classes. Overall, our results demonstrate that TabEBM consistently improves performance across different datasets and sample sizes, outperforming TabPFGen. Moreover, TabPFGen is limited in usability (e.g., it supports only up to ten classes), while TabEBM scales to any number of classes.

In a broader context, some recent work attempts to adapt Large Language Models (LLMs) for tabular data generation [25, 71, 9]. However, data contamination is an inherent issue with such LLM-based models [19, 36, 18, 49]. As the pre-training data is not typically open-source, these models can have unfair advantages in downstream tasks (i.e., the full real dataset, including the real test data, may have been used for pre-training). Therefore, in this paper, we focus on models without support from LLMs, thus avoiding potential biases from data contamination.

Data Augmentation (DA) for Tabular Data. DA is an omnipresent technique in computer vision and natural language processing [82, 73, 72, 60, 26, 2]. However, DA for tabular data remains

Methods

Category

Training

Generation

Practicability

Learned

distribution

Training-free

Class-specific

models

Stratified

generation

Unlimited

classes

ACC improve

(

10 classes)

ACC improve

(

>

10 classes)

SMOTE [13]

Interpolation

N/A

"

N/A

"

"

%

%

TVAE [85]

VAE

p

(

x

, y

)

%

%

%

"

%

%

CTGAN [85]

GAN

p

(

x

|

y

)

%

%

"

"

%

%

NFLOW [24]

Normal. Flows

p

(

x

, y

)

%

%

%

"

%

%

TabDDPM [42]

Diffusion

p

(

x

|

y

)

%

%

"

"

%

%

ARF [83]

Random Forest

p

(

x

, y

)

%

%

%

"

%

%

GOGGLE [47]

GNN

p

(

x

|

y

)

%

%

"

"

%

%

TabPFGen [48]

PFN

p

(

x

|

y

)

"

%

"

%

"

%

TabEBM (Ours)

PFN

p

(

x

y

)

"

"

"

"

"

"

|

underexplored, and existing methods often perform poorly in real-world tasks, sometimes even reduce performance [51]. Recent studies show that using the same transformations across all classes leads to varied performance impacts [3, 41], indicating that data augmentation effects are class-specific and suggesting that different classes may require distinct augmentations. Given the lack of symmetries in tabular data, we believe this class-dependent effect is even more pronounced. Therefore, we propose TabEBM as a class-specific generative model to produce tailored augmentations for each class.

Prior-fitted Networks (PFNs) for Tabular Data. Recent work proposes to approximate the posterior predictive distribution with transformers [59, 33, 61, 79, 20]. PFNs can be adapted for various purposes by pre-training the transformer with corresponding 'prior data', and then it can make in-context predictions with unseen downstream data. For instance, TabPFN is a variant that is pre-trained on a prior designed for tabular data [33]. We note that prior data is different to synthetic data in this paper. Specifically, prior data refers to manually crafted fake data (e.g., y = 2 x ) with no real-world semantics. In contrast, synthetic data from generators is expected to have the same semantics as real data. Inspired by TabPFN's success in small-size classification tasks, TabEBM converts TabPFN into multiple EBMs that learn the marginal distribution for each class. The training-free nature of TabPFN enables TabEBM to generate high-quality tabular data without introducing extra training costs. Additionally, our class-specific design lets TabEBM surpass TabPFN's limits and scale to more than ten classes.

Limitations and Future Work. TabEBM is a general method that relies on an underlying binary classifier, and as such, its strengths and weaknesses are directly tied to this classifier. We used TabPFN because it is a well-established open-source pre-trained model for tabular data. Therefore, TabEBM inherits some of TabPFN's limitations, particularly in scaling to a larger number of features. TabEBM can handle datasets with over 1000 samples, overcoming TabPFN's limitation, as it processes one class at a time. In Appendix D.5.3, we show that TabEBM outperforms other generators on larger datasets, though the performance gains decrease as the sample size increases. Although we implement TabEBM with TabPFN in this paper, we stress that TabEBM is compatible with any classifier that can be adapted into EBMs, as described in Section 2. As foundational models for tabular data evolve [81], new models capable of handling more features and samples are expected. Integrating them into TabEBM will enhance its ability to manage high-dimensional datasets, increasing its versatility and utility. Finally, note that, generators that are limited in modelling multivariate distributions may still perform well on univariate fidelity metrics, which is a standard approach to evaluating such models. However, evaluating their ability to learn more complex, high-order, relationships between features remains an open research question [78], which we leave for future work.

5 Conclusion

We introduced TabEBM, the first tabular data augmentation method that creates class-specific EBM generators, learning the marginal distribution for each class separately. We also provide the first comprehensive analysis of tabular data augmentation across various dataset sizes. Our results demonstrate that TabEBM improves downstream performance through data augmentation on realworld datasets, outperforming other benchmark generators. The statistical evaluation confirms that TabEBM generates high-fidelity synthetic data, particularly for small datasets. We release our method as an open-source library, allowing users to generate data immediately without additional training.

Acknowledgments and Disclosure of Funding

The authors would like to thank Francisco Vargas, Randall Balestriero, and Otilia Stretcu for their insightful discussions and valuable input early in the project. NS and MJ acknowledge the support of the U.S. Army Medical Research and Development Command of the Department of Defense; through the FY22 Breast Cancer Research Program of the Congressionally Directed Medical Research Programs, Clinical Research Extension Award GRANT13769713. Opinions, interpretations, conclusions, and recommendations are those of the authors and are not necessarily endorsed by the Department of Defense.

References

Hassane Alami, Lysanne Rivard, Pascale Lehoux, Steven J Hoffman, Stephanie Bernadette Mafalda Cadeddu, Mathilde Savoldelli, Mamane Abdoulaye Samri, Mohamed Ali Ag Ahmed, Richard Fleet, and Jean-Paul Fortin. Artificial intelligence in health care: laying the foundation for responsible, sustainable, and inclusive innovation in low-and middle-income countries. Globalization and Health , 16:1-6, 2020.

Antreas Antoniou, Amos Storkey, and Harrison Edwards. Data augmentation generative adversarial networks. arXiv preprint arXiv:1711.04340 , 2017.

Randall Balestriero, Leon Bottou, and Yann LeCun. The effects of regularization and data augmentation are class dependent. Advances in Neural Information Processing Systems , 35:3787837891, 2022.

Ms Aayushi Bansal, Dr Rewa Sharma, and Dr Mamta Kathuria. A systematic review on data scarcity problem in deep learning: solution and applications. ACM Computing Surveys (CSUR) , 54(10s):1-29, 2022.

Andreas D Baxevanis, Gary D Bader, and David S Wishart. Bioinformatics . John Wiley & Sons, 2020.

Anton Bespalov, Thomas Steckler, Bruce Altevogt, Elena Koustova, Phil Skolnick, Daniel Deaver, Mark J Millan, Jesper F Bastlund, Dario Doller, Jeffrey Witkin, et al. Failed trials for central nervous system disorders do not necessarily invalidate preclinical models and drug targets. Nature Reviews Drug Discovery , 15(7):516-516, 2016.

B. Bischl, Giuseppe Casalicchio, Matthias Feurer, Frank Hutter, Michel Lang, Rafael Gomes Mantovani, Jan N. van Rijn, and Joaquin Vanschoren. Openml benchmarking suites. Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks (NeurIPS 2021) , 2021.

Vadim Borisov, Tobias Leemann, Kathrin Seßler, Johannes Haug, Martin Pawelczyk, and Gjergji Kasneci. Deep neural networks and tabular data: A survey. IEEE Transactions on Neural Networks and Learning Systems , 2022.

Vadim Borisov, Kathrin Sessler, Tobias Leemann, Martin Pawelczyk, and Gjergji Kasneci. Language models are realistic tabular data generators. In The Eleventh International Conference on Learning Representations , 2022.

Leo Breiman. Random forests. Machine learning , 45(1):5-32, 2001.

Chenjing Cai, Shiwei Wang, Youjun Xu, Weilin Zhang, Ke Tang, Qi Ouyang, Luhua Lai, and Jianfeng Pei. Transfer learning for drug discovery. Journal of Medicinal Chemistry , 63(16):8683-8694, 2020.

Rees Chang, Yu-Xiong Wang, and Elif Ertekin. Towards overcoming data scarcity in materials science: unifying models and datasets with a mixture of experts framework. npj Computational Materials , 8(1):242, 2022.

Nitesh V Chawla, Kevin W Bowyer, Lawrence O Hall, and W Philip Kegelmeyer. Smote: synthetic minority over-sampling technique. Journal of artificial intelligence research , 16:321357, 2002.

Tianqi Chen and Carlos Guestrin. Xgboost: A scalable tree boosting system. In Proceedings of the 22nd acm sigkdd international conference on knowledge discovery and data mining , pages 785-794, 2016.

Tadeusz Ciecierski-Holmes, Ritvij Singh, Miriam Axt, Stephan Brenner, and Sandra Barteit. Artificial intelligence for strengthening healthcare systems in low-and middle-income countries: a systematic scoping review. npj Digital Medicine , 5(1):162, 2022.

David R Cox. The regression analysis of binary sequences. Journal of the Royal Statistical Society Series B: Statistical Methodology , 20(2):215-232, 1958.

Imre Csiszár. I-divergence geometry of probability distributions and minimization problems. The annals of probability , pages 146-158, 1975.

Jasper Dekoninck, Mark Niklas Müller, Maximilian Baader, Marc Fischer, and Martin Vechev. Evading data contamination detection for language models is (too) easy. arXiv preprint arXiv:2402.02823 , 2024.

Chunyuan Deng, Yilun Zhao, Xiangru Tang, Mark B. Gerstein, and Arman Cohan. Investigating data contamination in modern benchmarks for large language models. In North American Chapter of the Association for Computational Linguistics , 2024.

Samuel Dooley, Gurnoor Singh Khurana, Chirag Mohapatra, Siddartha V Naidu, and Colin White. Forecastpfn: Synthetically-trained zero-shot forecasting. Advances in Neural Information Processing Systems , 36, 2024.

Georgios Douzas and Fernando Bacao. Effective data generation for imbalanced learning using conditional generative adversarial networks. Expert Systems with applications , 91:464-471, 2018.

Dheeru Dua and Casey Graff. Uci machine learning repository, 2017.

Larry A Dunning and Ray Kresman. Privacy preserving data sharing with anonymous id assignment. IEEE transactions on information forensics and security , 8(2):402-413, 2012.

Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios. Neural spline flows. Advances in neural information processing systems , 32, 2019.

Xi Fang, Weijie Xu, Fiona Anting Tan, Ziqing Hu, Jiani Zhang, Yanjun Qi, Srinivasan H. Sengamedu, and Christos Faloutsos. Large language models (llms) on tabular data: Prediction, generation, and understanding - a survey. Transactions on Machine Learning Research , 2024.

Steven Y Feng, Varun Gangal, Jason Wei, Sarath Chandar, Soroush Vosoughi, Teruko Mitamura, and Eduard Hovy. A survey of data augmentation approaches for nlp. In Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021 , pages 968-988, 2021.

Evelyn Fix. Discriminatory analysis: nonparametric discrimination, consistency properties , volume 1. USAF school of Aviation Medicine, 1985.

Yury Gorishniy, Ivan Rubachev, Valentin Khrulkov, and Artem Babenko. Revisiting deep learning models for tabular data. Advances in Neural Information Processing Systems , 34:1893218943, 2021.

Will Grathwohl, Kuan-Chieh Wang, Joern-Henrik Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky. Your classifier is secretly an energy based model and you should treat it like one. In International Conference on Learning Representations , 2019.

Léo Grinsztajn, Edouard Oyallon, and Gaël Varoquaux. Why do tree-based models still outperform deep learning on typical tabular data? Advances in neural information processing systems , 35:507-520, 2022.

Lasse Hansen, Nabeel Seedat, Mihaela van der Schaar, and Andrija Petrovic. Reimagining synthetic tabular data generation through data-centric ai: A comprehensive benchmark. Advances in Neural Information Processing Systems , 36:33781-33823, 2023.

Mikel Hernandez, Gorka Epelde, Ane Alberdi, Rodrigo Cilla, and Debbie Rankin. Synthetic data generation for tabular health records: A systematic review. Neurocomputing , 493:28-45, 2022.

Noah Hollmann, Samuel Müller, Katharina Eggensperger, and Frank Hutter. TabPFN: A transformer that solves small tabular classification problems in a second. In The Eleventh International Conference on Learning Representations , 2023.

J. D. Hunter. Matplotlib: A 2d graphics environment. Computing in Science & Engineering , 9(3):90-95, 2007.

Dipendra Jha, Kamal Choudhary, Francesca Tavazza, Wei-keng Liao, Alok Choudhary, Carelyn Campbell, and Ankit Agrawal. Enhancing materials property prediction by leveraging computational and experimental data using deep transfer learning. Nature communications , 10(1):5316, 2019.

Minhao Jiang, Ken Ziyu Liu, Ming Zhong, Rylan Schaeffer, Siru Ouyang, Jiawei Han, and Sanmi Koyejo. Investigating data contamination for pre-training language models. arXiv preprint arXiv:2401.06059 , 2024.

Xiangjian Jiang, Andrei Margeloiu, Nikola Simidjievski, and Mateja Jamnik. Protogate: Prototype-based neural networks with global-to-local feature selection for tabular biomedical data. In Proceedings of the 41st International Conference on Machine Learning (ICML) , 2024.

Hao Jin, Yan Luo, Peilong Li, and Jomol Mathew. A review of secure and privacy-preserving medical data sharing. IEEE access , 7:61656-61669, 2019.

Marvin Karson. Handbook of methods of applied statistics. volume i: Techniques of computation descriptive methods, and statistical inference. volume ii: Planning of surveys and experiments. im chakravarti, rg laha, and j. roy, new york, john wiley; 1967., 1968.

Jayoung Kim, Chaejeong Lee, and Noseong Park. Stasy: Score-based tabular data synthesis. In The Eleventh International Conference on Learning Representations , 2022.

Polina Kirichenko, Mark Ibrahim, Randall Balestriero, Diane Bouchacourt, Shanmukha Ramakrishna Vedantam, Hamed Firooz, and Andrew G Wilson. Understanding the detrimental class-level effects of data augmentation. Advances in Neural Information Processing Systems , 36, 2024.

Akim Kotelnikov, Dmitry Baranchuk, Ivan Rubachev, and Artem Babenko. Tabddpm: Modelling tabular data with diffusion models. In International Conference on Machine Learning , pages 17564-17579. PMLR, 2023.

Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang. A tutorial on energybased learning. Predicting structured data , 1(0), 2006.

Chaejeong Lee, Jayoung Kim, and Noseong Park. Codi: Co-evolving contrastive diffusion models for mixed-type tabular synthesis. In International Conference on Machine Learning , pages 18940-18956. PMLR, 2023.

Guillaume Lemaître, Fernando Nogueira, and Christos K. Aridas. Imbalanced-learn: A python toolbox to tackle the curse of imbalanced datasets in machine learning. Journal of Machine Learning Research , 18(17):1-5, 2017.

Roman Levin, Valeriia Cherepanova, Avi Schwarzschild, Arpit Bansal, C Bayan Bruss, Tom Goldstein, Andrew Gordon Wilson, and Micah Goldblum. Transfer learning with deep tabular models. In The Eleventh International Conference on Learning Representations , 2022.

Tennison Liu, Zhaozhi Qian, Jeroen Berrevoets, and Mihaela van der Schaar. Goggle: Generative modelling for tabular data by learning relational structure. In The Eleventh International Conference on Learning Representations , 2023.

Junwei Ma, Apoorv Dankar, George Stein, Guangwei Yu, and Anthony Caterini. Tabpfgentabular data generation with tabpfn. In NeurIPS 2023 Second Table Representation Learning Workshop , 2023.

Inbal Magar and Roy Schwartz. Data contamination: From memorization to exploitation. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers) , pages 157-165, 2022.

Bradley A Malin, Khaled El Emam, and Christine M O'Keefe. Biomedical data privacy: problems, perspectives, and recent advances. Journal of the American medical informatics association , 20(1):2-6, 2013.

Dionysis Manousakas and Sergül Aydöre. On the usefulness of synthetic tabular data generation. In Data-centric Machine Learning Research (DMLR) Workshop at the 40th International Conference on Machine Learning (ICML) , 2023.

Andrei Margeloiu, Nikola Simidjievski, Pietro Lio, and Mateja Jamnik. Weight predictor network with feature selection for small sample tabular biomedical data. AAAI Conference on Artificial Intelligence , 2023.

Calvin McCarter. What exactly has tabpfn learned to do? In ICLR Blogposts 2024 , 2024.

Duncan McElfresh, Sujay Khandagale, Jonathan Valverde, Vishak Prasad C, Ganesh Ramakrishnan, Micah Goldblum, and Colin White. When do neural nets outperform boosted trees on tabular data? Advances in Neural Information Processing Systems , 36, 2024.

Mary L McHugh. The chi-square test of independence. Biochemia medica , 23(2):143-149, 2013.

Daniele Micci-Barreca. A preprocessing scheme for high-cardinality categorical attributes in classification and prediction problems. ACM SIGKDD explorations newsletter , 3(1):27-32, 2001.

Daniel J Mollura, Melissa P Culp, Erica Pollack, Gillian Battino, John R Scheel, Victoria L Mango, Ameena Elahi, Alan Schweitzer, and Farouk Dako. Artificial intelligence in low-and middle-income countries: innovating global health radiology. Radiology , 297(3):513-520, 2020.

LaRonda L Morford, Christopher J Bowman, Diann L Blanset, Ingrid B Bøgh, Gary J Chellman, Wendy G Halpern, Gerhard F Weinbauer, and Timothy P Coogan. Preclinical safety evaluations supporting pediatric drug development with biopharmaceuticals: strategy, challenges, current practices. Birth Defects Research Part B: Developmental and Reproductive Toxicology , 92(4):359-380, 2011.

Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Transformers can do bayesian inference. In International Conference on Learning Representations , 2022.

Alhassan Mumuni and Fuseini Mumuni. Data augmentation: A comprehensive survey of modern approaches. Array , 16:100258, 2022.

Thomas Nagler. Statistical foundations of prior-data fitted networks. In International Conference on Machine Learning , pages 25660-25676. PMLR, 2023.

Mehmet Ercan Nergiz, Maurizio Atzori, and Christopher W Clifton. δ -presence. Encyclopedia of Cryptography, Security and Privacy , pages 1-5, 2019.

Noseong Park, Mahmoud Mohammadi, Kshitij Gorde, Sushil Jajodia, Hongkyu Park, and Youngmin Kim. Data synthesis based on generative adversarial networks. Proceedings of the VLDB Endowment , 11(10), 2018.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems , 32, 2019.

F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research , 12:2825-2830, 2011.

Liudmila Prokhorenkova, Gleb Gusev, Aleksandr Vorobev, Anna Veronika Dorogush, and Andrey Gulin. Catboost: unbiased boosting with categorical features. Advances in neural information processing systems , 31, 2018.

Zhaozhi Qian, Rob Davis, and Mihaela van der Schaar. Synthcity: a benchmark framework for diverse use cases of tabular synthetic data. Advances in Neural Information Processing Systems , 36, 2024.

Yiming Qin, Huangjie Zheng, Jiangchao Yao, Mingyuan Zhou, and Ya Zhang. Class-balancing diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 18434-18443, 2023.

Sebastian Raschka. Model evaluation, model selection, and algorithm selection in machine learning. arXiv preprint arXiv:1811.12808 , 2018.

Vignesh Sampath, Iñaki Maurtua, Juan Jose Aguilar Martin, and Aitor Gutierrez. A survey on generative adversarial networks for imbalance problems in computer vision tasks. Journal of big Data , 8:1-59, 2021.

Nabeel Seedat, Nicolas Huynh, Boris van Breugel, and Mihaela van der Schaar. Curated llm: Synergy of llms and data curation for tabular augmentation in low-data regimes. In Forty-first International Conference on Machine Learning , 2024.

Connor Shorten and Taghi M Khoshgoftaar. A survey on image data augmentation for deep learning. Journal of big data , 6(1):1-48, 2019.

Connor Shorten, Taghi M Khoshgoftaar, and Borko Furht. Text data augmentation for deep learning. Journal of big Data , 8(1):101, 2021.

Ravid Shwartz-Ziv and Amitai Armon. Tabular data: Deep learning is not all you need. Information Fusion , 81:84-90, 2022.

Theresa Stadler and Carmela Troncoso. Why the search for a privacy-preserving data sharing mechanism is failing. Nature Computational Science , 2(4):208-210, 2022.

Fahim Sufi. Addressing data scarcity in the medical domain: A gpt-based approach for synthetic data generation and feature extraction. Information , 15(5):264, 2024.

Haoyuan Sun, Navid Azizan, Akash Srivastava, and Hao Wang. Private synthetic data meets ensemble learning. arXiv preprint arXiv:2310.09729 , 2023.

Ruibo Tu, Zineb Senane, Lele Cao, Cheng Zhang, Hedvig Kjellström, and Gustav Eje Henter. Causality for tabular data synthesis: A high-order structure causal benchmark framework. arXiv preprint arXiv:2406.08311 , 2024.

Jordan Ubbens, Ian Stavness, and Andrew G Sharpe. Gpfn: Prior-data fitted networks for genomic prediction. bioRxiv , pages 2023-09, 2023.

Boris Van Breugel, Zhaozhi Qian, and Mihaela Van Der Schaar. Synthetic data, real errors: how (not) to publish and use synthetic data. In International Conference on Machine Learning , pages 34793-34808. PMLR, 2023.

Boris van Breugel and Mihaela van der Schaar. Why tabular foundation models should be a research priority. In Forty-first International Conference on Machine Learning , 2024.

David A Van Dyk and Xiao-Li Meng. The art of data augmentation. Journal of Computational and Graphical Statistics , 10(1):1-50, 2001.

David S Watson, Kristin Blesch, Jan Kapar, and Marvin N Wright. Adversarial random forests for density estimation and generative modeling. In International Conference on Artificial Intelligence and Statistics , pages 5357-5375. PMLR, 2023.

Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11) , pages 681-688. Citeseer, 2011.

Lei Xu, Maria Skoularidou, Alfredo Cuesta-Infante, and Kalyan Veeramachaneni. Modeling tabular data using conditional gan. Advances in neural information processing systems , 32, 2019.

Aiqing Zhang and Xiaodong Lin. Towards secure and privacy-preserving data sharing in e-health systems via consortium blockchain. Journal of medical systems , 42(8):140, 2018.

Hengrui Zhang, Jiani Zhang, Zhengyuan Shen, Balasubramaniam Srinivasan, Xiao Qin, Christos Faloutsos, Huzefa Rangwala, and George Karypis. Mixed-type tabular data synthesis with score-based diffusion in latent space. In The Twelfth International Conference on Learning Representations , 2023.

Zilong Zhao, Aditya Kunar, Robert Birke, and Lydia Y Chen. Ctab-gan: Effective table data synthesizing. In Asian Conference on Machine Learning , pages 97-112. PMLR, 2021.

Xu Zheng and Zhipeng Cai. Privacy-preserved data sharing towards multiple parties in industrial iots. IEEE Journal on Selected Areas in Communications , 38(5):968-979, 2020.

Appendix

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

Table of Contents

A Broader Impact Statement

This paper introduces a novel data augmentation approach, TabEBM, that aims to advance the field of machine learning by addressing challenges in the low-sample-size regime. Furthermore, TabEBM offers an elegant solution to learning the unique features in generating samples for each class, leading to high-fidelity synthetic data that can effectively improve downstream performance. These characteristics can be particularly useful in data-scarce domains like healthcare (e.g., pre-clinical drug evaluation in early-stage clinical trials [6, 58]). Moreover, we also demonstrate that TabEBM is readily applicable for privacy-preserving data sharing in high-stake tasks [89, 77].

TabEBM's impact further extends to enabling broader machine learning applications in data-scarce domains, for instance, facilitating data analysis in clinical scenarios with limited access to data collection techniques. Improving the performance of machine learning models in such applications can further foster the uptake of more sophisticated ML approaches and, ultimately, help improve the quality of healthcare [1, 15, 57]. TabEBM can further facilitate research and enhance machine learning accessibility in various communities across societal and scientific domains. To this end, our work has only been evaluated in a strictly research setting. Further applications of our work in scenarios with sensitive data bear some risks. As TabEBM is a generative model, training models with the resulting generated samples can bias the downstream model. Therefore, this risk, together with other data privacy risks during downstream deployment, must be carefully managed.

B Reproducibility

B.1 Datasets

All eight datasets are publicly available on OpenML [7], and their details are listed in Table 3. To ensure consistent stratified data-splitting across all datasets, we remove classes with fewer than 10 samples. For example, the original 'energy' dataset contains 14 classes with fewer than 10 samples, which could result in a validation set lacking samples from these classes, leading to unstratified data splitting.

B.2 Data Splitting

Figure 9 shows the data splitting setup used across all datasets. Note that data sharing (Section 3.3) shares the same data splitting as data augmentation, except that the 'Training set' and 'Validation set' containing real data are no longer used for training the downstream predictors.

Dataset

𝑁

Oracle set (i.e., full train set)

test

Test set

𝑁

test

= min

𝑁

2

, 500

Subsets

𝑁 real

∈ [20, 50, 100, 200, 500]

Training set

𝑁 real 	×	0.8

Validation set

𝑁 real 	×	0.2

Generator

Downstream

predictors

Synthetic data

500

⊕

Predictions

B.3 Data Preprocessing

Following the procedures presented in prior work [54, 30], we perform preprocessing in two steps. We first compute the required statistics with training data and then transform it. Firstly, we impute the missing values with the mean value for numerical features and the most mode value for categorical features. Secondly, we convert the categorical features into numerical features equal to Leave-one-out Target Statistic [66, 56]. Next, we perform Z-score normalisation for each feature. Specifically, we compute each feature's mean and standard deviation in the training data and then transform the training samples to have a mean of zero and a variance of one for each feature. Finally, we apply the same transformation to the validation and test data before conducting evaluations.

B.4 Software and Computing Resources

Software implementation. (i) For generators: We implemented TabEBM using PyTorch 1.13 [64], an open-source deep learning library with a BSD licence. We implemented SMOTE with Imbalancedlearn [45], an open-source Python library for imbalanced datasets with an MIT licence. For other benchmark generators, we used their open-source implementations in Synthcity [67], a library for generating and evaluating synthetic tabular data with an Apache-2.0 license. (ii) For downstream predictors: We implemented TabPFN with its open-source implementation ( https://github.com/ automl/TabPFN ). We implemented the other five downstream predictors (i.e., Logistic Regression, KNN, MLP, Random Forest and XGBoost) with their open-source implementation in scikit-learn [65], an open-source Python library under the 3-Clause BSD license. (iii) For result analysis and visualisation: All numerical plots and graphics have been generated using Matplotlib 3.7 [34], a Pythonbased plotting library with a BSD licence. The model architecture was generated using draw.io ( https://github.com/jgraph/drawio ), a free drawing software under Apache License 2.0.

We ensure the consistency and reproducibility of experimental results by implementing a uniform pipeline using PyTorch Lightning, an open-source library under an Apache-2.0 licence. We further fixed the random seeds for data loading and evaluation throughout the training and evaluation process. This ensured that TabEBM and all benchmark models were trained and evaluated on the same set of samples. The experimental environment settings, including library dependencies, are specified in the open-source library for reference and reproduction purposes.

Computing Resources. We trained 140,000 models for evaluations (including over 35,000 of generators and over 10,500 for downstream predictors). All our experiments are run on a single machine from an internal cluster with a GPU Nvidia Quadro RTX 8000 with 48GB memory and an Intel(R) Xeon(R) Gold 5218 CPU with 16 cores (at 2.30GHz). The operating system was Ubuntu 20.4.4 LTS.

B.5 TabEBM open-source library

We implemented TabEBM as an extensible library, and the code is available on https://github.com/andreimargeloiu/TabEBM . For practitioners, it offers an easy-to-use, domain-agnostic tool that requires no training, making it particularly suitable for data augmentation, especially in small datasets. For researchers, the library includes the complete implementation of TabEBM, facilitating future extensions and investigations into class-specific energy-based models.

The library has two core functionalities:

Generate synthetic data : The library can generate data for augmentation.

Compute and visualise the energy function : The library allows computation of TabEBM's energy function and the unnormalised data density. The demo notebook, TabEBM_approximated_density.ipynb , shows the TabEBM-inferred densities under conditions of data noise and class imbalance (thus recreating the plots from Appendix C).

B.6 Implementation of Generators

TabEBM. In all our experiments, the surrogate binary classifier in TabEBM is a pretrained in-context model, TabPFN [33], using the official model weights released by the authors ( https://github.com/automl/TabPFN/raw/main/tabpfn/models_diff/ prior_diff_real_checkpoint_n_0_epoch_42.cpkt ). We use TabPFN with three ensembles. We use four surrogate negative samples, X neg c , positioned at α neg dist = 5 standard deviations from zero, in random corners of a hypercube in R D (as explained in Section 2.2), distant from any real data. In Appendix D.1, we show that TabEBM is robust to the distribution of the negative samples.

We use SGLD [84] for sampling from TabEBM, where the starting points x synth 0 are initialised by adding Gaussian noise with zero mean and standard deviation σ start = 0 . 01 to a randomly selected sample of the specific class, i.e., x synth 0 ∼ N ( X c , σ 2 start I ) . For SGLD, we used the following parameters: step size α step = 0 . 1 , noise scale α noise = 0 . 01 and number of steps T = 200 . We found TabEBM to be robust to the SGLD settings (see Appendix D.2).

TabPFGen. We re-implemented TabPFGen [48] by closely following the original paper since no official implementation is available. As recommended in [48], the starting points are initialised by adding Gaussian noise with zero mean and standard deviation of 0.01 to the training points.

SMOTE. We use the open-source implementation of SMOTE from Imbalanced-learn [45], and the number neighbours k is set within the range of { 1 , 3 , 5 } . When applicable, we always set the maximum value for nearest neighbours (i.e., k = 5 ). However, very low-sample-size datasets may not contain sufficient samples for large k . For instance, the 'fourier' dataset ( N real = 20 ) only has two samples per class. We set k = 1 to generate synthetic data with SMOTE in these cases.

For the other six benchmark generators, we use their open-source implementations in Synthcity [67]. Following prior studies [87, 80, 71, 48], we use the default settings for all generators.

B.7 Implementation of Downstream Predictors

We implemented TabPFN with its official implementation [33] and the other five downstream predictors with the scikit-learn library [65]. Following prior studies [80, 71], we use the default settings for all downstream predictors.

C Limitations of Existing Generative Methods

We showcase three limitations of current generative models: (1) Appendix C shows that models approximating the joint distribution p ( x , y ) may fail to preserve the stratification of the real data and even fail to generate samples from specific classes. (2) Appendix C evaluates the approximated class-conditional distributions p ( x | y ) on data with increasing noise levels, and (3) Appendix C evaluates the approximated class-conditional distributions p ( x | y ) on data with increasing class imbalance.

1

2

3

4

5

6

7

8

9

10 11 12 13 14 15 16 17 18 19 20 21 22 23

Class ID

0

5

10

15

20

Class Distribution (%)

Real Data

TVAE

The approximated (unnormormalized) probability p(x | y). circles

Real data

TabPFGen inferred p(x|y=0)

TabPFGen inferred p(x|y=1)

TabEBM (ours) inferred p(x|y=0)

TabEBM (ours) inferred p(x|y=1)

(a) Noise level 0.1

Real data

TabPFGen inferred p(x|y=0)

TabPFGen inferred p(x|y=1)

TabEBM (ours) inferred p(x|y=0)

TabEBM (ours) inferred p(x|y=1)

The approximated (unnormormalized) probability p(x | y). circles

(b) Noise level 0.25

Real data

TabPFGen inferred p(x|y=0)

TabPFGen inferred p(x|y=1)

TabEBM (ours) inferred p(x|y=0)

TabEBM (ours) inferred p(x|y=1)

The approximated (unnormormalized) probability p(x | y). circles

(c) Noise level 0.5

Real data

TabPFGen inferred p(x|y=0)

TabPFGen inferred p(x|y=1)

TabEBM (ours) inferred p(x|y=0)

TabEBM (ours) inferred p(x|y=1)

The approximated (unnormormalized) probability p(x | y). circles

(d) Noise level 1

Real data

TabPFGen inferred p(x|y=0)

TabPFGen inferred p(x|y=1)

TabEBM (ours) inferred p(x|y=0)

TabEBM (ours) inferred p(x|y=1)

The approximated (unnormormalized) probability p(x | y). circles

(e) Noise level 2

The approximated (unnormormalized) probability p(x | y). two_blobs

Real data

TabPFGen inferred p(x|y=0)

TabPFGen inferred p(x|y=1)

TabEBM (ours) inferred p(x|y=0)

TabEBM (ours) inferred p(x|y=1)

(a) Class ratio 150:150

Real data

TabPFGen inferred p(x|y=0)

TabPFGen inferred p(x|y=1)

TabEBM (ours) inferred p(x|y=0)

TabEBM (ours) inferred p(x|y=1)

The approximated (unnormormalized) probability p(x | y). two_blobs

(b) Class ratio 50:250

Real data

TabPFGen inferred p(x|y=0)

TabPFGen inferred p(x|y=1)

TabEBM (ours) inferred p(x|y=0)

TabEBM (ours) inferred p(x|y=1)

The approximated (unnormormalized) probability p(x | y). two_blobs

(c) Class ratio 25:275

Real data

TabPFGen inferred p(x|y=0)

TabPFGen inferred p(x|y=1)

TabEBM (ours) inferred p(x|y=0)

TabEBM (ours) inferred p(x|y=1)

The approximated (unnormormalized) probability p(x | y). two_blobs

(d) Class ratio 10:290

D Extended Experimental Results

D.1 Ablations on the distribution of the surrogate negative samples

D.1.1 Ablations on placing the negative samples Energy E(x|y=0) of TabEBM on one class only

Real data for one class

real samples

c

negative samples

c

real samples

c

negative samples neg

c

real samples

c

negative samples neg

c

Appendix D.1.1 shows TabEBM's energy E c ( x ) when varying the selection of the negative samples. TabEBM infers an accurate energy surface with distant negative samples, and the energy surface becomes inaccurate when negative samples resemble real samples. This occurs because TabPFN is uncertain when points of different classes are close, affecting its logits magnitude and making them unsuitable for density estimation.

D.1.2 Varying the number of negative samples

We evaluate the impact of the ratio |X neg c | : |X c | between the negative samples X neg c and the real samples |X c | . We vary |X neg c | while keeping |X c | fixed, simulating both balanced and highly imbalanced scenarios. The negative samples are placed in random corners of the hypercube (as described in Section 2), at five standard deviations in each direction (i.e., α neg dist = 5 ). To ensure reliable outcomes, we maintained a consistent ratio across all classes, keeping the same proportion of negative samples for each class.

Table 4 shows the results across six datasets with N real = 100 real samples, demonstrating that TabEBM is robust to imbalances in the surrogate binary tasks. The column with |X neg c | = 4 represents the TabEBM results from the main paper, where four negative samples were placed in the corners (as described in Section 2). There are negligible differences in performance, and TabEBM consistently outperforms both the baseline and other generators (as shown in Table 1).

D.1.3 Varying the distance of the negative samples

We assess the effect of varying the distance of negative samples. We use TabEBM with four negative samples positioned randomly at the corners of the hypercube, as outlined in Section 2 (this corresponds to the experimental setup from the main paper). The distance of the negative samples, denoted as α neg dist , is varied. Table 5 demonstrates that TabEBM remains generally robust to changes in this distance, with only small performance variations across different datasets. Importantly, using TabEBM for data augmentation consistently improves performance by approximately 3% compared to the Baseline, regardless of the distance used.

D.2 Ablations on the sensitivity to the hyperparameters of SGLD sampling

We vary two key hyperparameters of SGLD on the 'biodeg' binary dataset with N real = 100 : the step size α step and the noise scale α noise. Table 6 shows that TabEBM remains stable with respect to these hyperparameters. Note that smaller values of α noise are expected to perform better because SGLD sampling adds noise at each iteration (see Line 7 in Algorithm 1), thus larger values of α noise will hinder convergence of the SGLD sampler.

D.3 Distribution of Logits and Unnormalized Density in TabEBM

0

10

20

30

Distance from the real data

0.0

2.5

5.0

7.5

10.0

12.5

15.0

Logit value

biodeg

Logit

f

(

x

)[1] for real data

Logit

f

(

x

)[0] for negative samples

0

10

20

30

Distance from the real data

0.0

2.5

5.0

7.5

10.0

12.5

15.0

Logit value

fourier

Logit

f

(

x

)[1] for real data

Logit

f

(

x

)[0] for negative samples

0

10

20

30

Distance from the real data

0.0

2.5

5.0

7.5

10.0

12.5

15.0

Logit value

steel

Logit

f

(

x

)[1] for real data

Logit

f

(

x

)[0] for negative samples

0

10

20

30

Distance from the real data

0.0

2.5

5.0

7.5

10.0

12.5

15.0

Logit value

protein

Logit

f

(

x

)[1] for real data

Logit

f

(

x

)[0] for negative samples

0

10

20

30

Distance from the real data

0

100

200

300

400

500

Unnorm. density (x1000)

biodeg

TabEBM unnorm density

0

10

20

30

Distance from the real data

0

25

50

75

100

125

Unnorm. density (x1000)

fourier

TabEBM unnorm density

0

10

20

30

Distance from the real data

0

50

100

150

Unnorm. density (x1000)

steel

TabEBM unnorm density

0

10

20

30

Distance from the real data

0

20

40

60

80

100

Unnorm. density (x1000)

protein

TabEBM unnorm density

D.4 Complete Trade-off Figures with Error Bars

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

(a1)

Inverse KL (train, syn)

↑

-

20

-

10

0

10

Normalised accuracy improvement (%)

SMOTE

TVAE

CTGAN

NFLOW

TabDDPM

ARF

GOGGLE

TabPFGen

TabEBM (Ours)

0

.

6

0

.

7

0

.

8

0

.

9

1

.

0

(a2)

KS test (train, syn)

↑

-

20

-

10

0

10

SMOTE

TVAE

CTGAN

NFLOW

TabDDPM

ARF

GOGGLE

TabPFGen

TabEBM (Ours)

0

.

2

0

.

3

0

.

4

0

.

5

(b1)

DCR (train, syn)

-

40

-

20

0

10

Normalised accuracy change (%)

SMOTE

TVAE

CTGAN

NFLOW

TabDDPM

ARF

GOGGLE

TabPFGen

TabEBM (Ours)

-

0

.

5

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

(b2)

d

-presence (train, syn)

-

40

-

20

0

10

SMOTE

TVAE

CTGAN

NFLOW

TabDDPM

ARF

GOGGLE

TabPFGen

TabEBM (Ours)

↑

↓

D.5 Results on Data Augmentation

D.5.1 Results on eight OpenML datasets.

Datasets

N

real

Baseline

(Real data)

SMOTE

TVAE

CTGAN

NFLOW

TabDDPM

ARF

GOGGLE

TabPFGen

TabEBM

At most 10 classes

protein

20

27.80

±

4.37

N/A

19.21

±

3.80

20.58

±

4.63

20.80

±

4.34

18.89

±

4.37

23.97

±

3.05

10.55

±

1.61

33.42

±

5.95

34.63

±

5.78

50

55.24

±

3.46

59.85

±

3.87

43.58

±

6.20

37.37

±

8.01

34.42

±

6.65

21.70

±

7.43

46.02

±

2.60

13.54

±

4.22

57.63

±

2.82

58.88

±

3.99

100

74.31

±

3.49

80.05

±

3.16

68.15

±

5.54

71.10

±

2.55

57.89

±

6.13

59.28

±

8.23

65.84

±

3.03

23.69

±

10.40

77.60

±

4.03

78.26

±

3.75

200

88.67

±

1.53

91.79

±

1.42

87.05

±

2.85

86.69

±

2.85

83.29

±

2.42

87.39

±

2.95

85.49

±

2.49

77.63

±

6.11

90.77

±

1.37

90.94

±

1.46

500

97.31

±

0.69

97.69

±

0.77

97.51

±

0.85

97.58

±

0.85

96.89

±

0.62

97.44

±

0.85

97.40

±

0.60

97.35

±

0.61

97.24

±

0.80

97.28

±

0.62

fourier

20

30.06

±

6.85

N/A

22.00

±

4.62

20.10

±

4.31

14.52

±

3.96

12.22

±

2.40

21.64

±

5.91

14.64

±

3.94

N/A

36.56

±

4.96

50

53.62

±

4.71

53.08

±

3.34

45.82

±

4.29

37.46

±

5.82

28.78

±

2.78

22.74

±

5.11

42.14

±

3.09

11.30

±

1.50

53.15

±

3.50

53.82

±

3.92

100

64.62

±

4.14

63.66

±

3.92

56.68

±

3.02

54.78

±

2.80

45.50

±

4.50

49.36

±

8.51

54.74

±

2.78

21.40

±

4.29

65.95

±

3.49

65.40

±

3.61

200

71.62

±

2.59

70.56

±

3.61

66.48

±

3.82

66.14

±

4.02

62.64

±

2.60

72.12

±

2.64

65.04

±

3.20

52.18

±

7.35

69.93

±

3.91

72.48

±

3.08

500

77.66

±

1.61

77.50

±

1.08

76.80

±

1.34

77.82

±

1.24

73.90

±

1.76

79.16

±

2.05

75.70

±

2.11

74.36

±

2.53

77.30

±

0.42

77.40

±

1.28

biodeg

20

65.26

±

8.01

68.72

±

4.50

69.02

±

5.37

59.39

±

6.25

58.28

±

8.30

50.00

±

0.00

58.45

±

8.16

51.80

±

4.07

70.68

±

4.94

71.18

±

5.25

50

75.27

±

2.63

74.65

±

3.28

73.44

±

4.02

70.21

±

3.61

55.68

±

9.27

50.00

±

0.00

72.74

±

3.74

55.75

±

7.45

75.69

±

2.44

75.56

±

3.22

100

78.92

±

1.98

77.78

±

2.65

77.27

±

3.15

77.71

±

1.81

63.50

±

10.77

57.50

±

6.27

77.25

±

1.66

65.87

±

6.72

78.15

±

1.45

79.00

±

1.99

200

82.59

±

1.84

81.42

±

1.27

80.48

±

1.82

80.19

±

2.48

79.16

±

2.49

80.45

±

1.48

80.88

±

1.68

80.66

±

1.49

82.56

±

1.68

82.58

±

1.90

500

85.00

±

0.70

84.37

±

0.75

84.40

±

0.68

84.67

±

0.98

84.45

±

0.91

84.58

±

0.70

84.68

±

1.06

83.66

±

0.67

84.56

±

0.98

84.55

±

0.92

steel

20

56.77

±

4.17

55.95

±

4.30

56.03

±

4.37

55.62

±

4.80

52.52

±

4.64

50.00

±

0.00

52.39

±

3.13

50.05

±

0.17

64.80

±

5.66

65.87

±

6.14

50

82.34

±

8.38

63.42

±

3.93

62.08

±

2.69

63.98

±

4.08

52.92

±

4.72

50.64

±

2.01

61.32

±

4.55

50.36

±

1.09

84.70

±

7.84

86.30

±

6.73

100

97.37

±

1.37

73.06

±

4.46

71.96

±

5.40

72.23

±

4.15

56.34

±

6.30

80.87

±

20.44

69.29

±

5.70

51.18

±

3.24

97.49

±

1.21

97.81

±

1.49

200

98.84

±

0.70

82.32

±

2.88

81.78

±

3.36

83.24

±

2.68

82.92

±

6.21

99.35

±

0.70

86.40

±

4.22

64.42

±

11.35

98.80

±

0.73

98.96

±

0.71

500

99.74

±

0.29

94.27

±

2.39

94.93

±

1.89

96.98

±

1.34

98.32

±

1.19

99.88

±

0.15

95.70

±

1.50

98.56

±

0.52

99.77

±

0.30

99.74

±

0.29

stock

20

83.18

±

4.37

83.69

±

3.10

74.01

±

5.09

56.92

±

16.52

74.99

±

6.60

78.73

±

12.25

69.64

±

6.88

73.40

±

4.88

82.95

±

4.44

83.81

±

4.94

50

90.01

±

2.07

90.01

±

2.43

82.27

±

4.30

78.91

±

4.14

78.94

±

8.78

89.68

±

1.92

83.72

±

2.50

79.00

±

6.87

89.95

±

2.08

90.15

±

1.76

100

92.39

±

1.06

92.09

±

1.45

90.75

±

2.20

89.43

±

3.29

86.16

±

3.83

92.12

±

1.16

90.17

±

1.92

89.30

±

1.33

92.12

±

1.12

92.57

±

1.27

200

94.16

±

0.92

93.99

±

0.70

93.57

±

1.10

93.28

±

1.59

91.92

±

2.00

94.22

±

1.10

93.05

±

1.35

92.07

±

1.76

94.17

±

0.89

94.16

±

1.07

Average rank

3.08

±

1.22

4.23

±

2.32

6.12

±

1.57

6.29

±

2.07

8.42

±

1.32

6.12

±

3.38

6.54

±

1.61

8.83

±

1.46

3.12

±

1.80

2.23

±

1.83

D.5.2 Results on six UCI Datasets

D.5.3 Results on larger sample sizes

D.6 Results on Statistical Fidelity

We aim to provide a fair and coherent comparison between TabEBM and existing methods and thus we follow the widely adopted evaluation process in prior studies. Specifically, we compute the three statistical fidelity metrics with the open-source implementations from the well-established benchmark, Synthcity. We note that the previous studies [87, 67] often operate under the assumption that the issues associated with multiple comparisons are less pronounced in generating low-dimensional tabular data, hence correction methods for multiple hypothesis testing are seldom employed. Following such assumptions, correction methods are not employed in this work. In addition, we would like to point out the imperfection of widely adopted univariate metrics (i.e., Inverse KL, KS test and test) in existing work. However, evaluating generators' ability to capture the joining feature relationships remains an open research question [78]. We leave this for future work to explore.

D.6.1 Similarity between Real Train Data and Synthetic Data

D.6.2 Similarity between Real Test Data and Synthetic Data

D.7 Results on Privacy Preservation

D.7.1 Downstream Accuracy in Data Sharing

D.7.2 DCR Evaluation

D.7.3 Delta-presence Evaluation

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Section 1 details our research objectives and highlights our contributions.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Presented in Section 4.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Section 2 presents the theoretical results of our proposed method.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Refer to Appendix B, where we provide full details on reproducing the results in the paper. We provide an open-source library of the proposed method.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Refer to Appendix B. All datasets used in this paper are publicly available, and the implementations of benchmark generators are open-source. We also provide an open-source library https://github.com/andreimargeloiu/TabEBM .

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Appendix B provides full descriptions of the experimental setup.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Refer to Section 3, where we provide standard deviations for all tables. Figure 6 and Figure 4 (Right) contain error bars. Due to the page limit, the error bars for all other figures are available in Appendix D. In Section 3.2 and Appendix D.6, we show statistical significance tests of the similarity between real data and synthetic data.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Refer to Appendix B.4, where we provide full details on the computation resources used in the paper.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We carefully check the NeurIPS Code of Ethics, and we confirm that our work follows the Code in every respect.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: Refer to Appendix A, where we include the societal impacts of our work.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: [NA]

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: Refer to Appendix B, where we provide the open-source licenses followed by the creators or original owners of assets.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We provide the implementation of our method as a python library attached to this submission. We will make it publicly available post-publication.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: [NA]

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [No]

Justification:[NA]

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors address the problem of data augmentation for classification tasks by proposing a novel approach that utilizes separate energy-based models (EBMs) to generate synthetic data for each class. These EBMs are derived directly from the logits of a binary classifier whose goal is to distinguish real data from negative data. These logits (their logsumexp over positive and negative classes) are reinterpreted them as an energy function for each class.\nIn order to avoid training these classifiers on the data itself, the authors take advantage of a pre-trained Prior-Data Fitted Network (PFN).\n\nThe authors perform a comprehensive evaluation of their method comparing it to competing methods for data augmentation on a number of small real world datasets. They show that their method is the only one that offers a consistent improvement in classification from the addition of synthetic data.\n\nSTRENGTHS:\n- The paper is well-structured and clearly written, effectively communicating the key ideas and concepts.\n- The authors introduce a novel approach by leveraging a Prior-Data Fitted Network (PFN) to obtain energy-based models (EBMs), which, to the best of my knowledge, has not been previously explored.\n- The experimental results section is thorough and comprehensive. They compare their approach to a wide range of relevant and competitive methods for synthetic data generation.\n- The appendices are well-organized and provide valuable supplementary information. The paper strikes a good balance with the most relevant information presented in main text and additional details provided in the appendix.\n\nWEAKNESSES:\n- Interpreting classifier logits as energy functions may not necessarily yield good density models. Given a classifier with logits $f(x)[y]$, one can add an arbitrary function $g(x)$ that is constant over $y$, $f(x)[y] \\rightarrow f(x)[y] + g(x)$, without affecting the conditional probabilities $p(y\\vert x)$ due to the invariance of the softmax function to the addition of a constant. However, this changes the energy function for the density model to $E(x) \\rightarrow E(x) - g(x)$. Consequently, there exists a functional degree of freedom that can arbitrarily alter the energy function without impacting the discriminative loss optimized by a classifier. As a result, a classifier trained solely to distinguish between positive and negative samples is not inherently encouraged to learn a useful energy function for the density over $x$.\nMy interpretation of the results of [1], from which the authors seem to draw inspiration is that reinterpreting the logits of a classification model as an energy based model by itself accomplishes nothing if the training process of the original classifier is not changed to encourage a good density model to be learned and encoded into the free degree of freedom. [1] achieves this by optimizing the joint likelihood, $E_{x,y} p_\\theta(y,x)$, instead of the conditional likelihood, $E_{x,y} p_\\theta(y\\vert x)$, used in training classifiers. This significantly complicates training by requiring sampling from the EBM during training. Therefore, it is unclear why the authors assume that the pre-trained PFN has learned useful energy functions for density modeling when it is derived from a simple classifier and this is not discussed in the paper.\n\n- Another weakness of the paper is in the experimental setup where only default hyperparameters are used for all competing methods. For many of these methods, tuning hyperparameters can substantially improve their results and the default hyperparameters may not be optimal.\nPersonally, I would find the presented results more convincing if the paper compared with fewer competing methods but made a better effort of getting good results out of those as it is impossible to know how much the authors tuned the degrees of freedom available during experimentation with their own proposed method.\n\n[1] Will Grathwohl, Kuan-Chieh Wang, Jörn-Henrik Jacobsen, David Duvenaud, Mohammad Norouzi, Kevin Swersky. Your Classifier is Secretly an Energy Based Model and You Should Treat it Like One, 2022 (https://arxiv.org/abs/1912.03263)\n\nQUESTIONS:\n- Can you provide a convincing explanation for why the logits from discriminative classifiers derived from TabPFN should be useful when reinterpreted as a class-conditional energy function in light of the first point above? Without a clear justification, it is difficult to understand the basis for expecting this method to work effectively.\n\n- Since the authors marginalize over $y$ by taking the LogSumExp over the $y=0$ and $y=1$ logits corresponding to negative and empirical class data, respectively, one would expect the energy function to represent a mixture distribution between these two densities. Why aren't the logits for $y=1$ used directly as an energy function instead? In other words, what is the motivation for marginalizing over $y$ when it is simply an indicator of whether the data is real or negative?\n\n- The details of how the negative examples are generated are not clear. The text simply states: \n\"These surrogate negative samples $\\mathcal{X}^{neg}_c$ are constructed to be far from the distribution of real data, ensuring that the classifier can easily distinguish them. This placement prevents class ambiguity and facilitates a robust energy landscape\". The appendix does provide some additional details, but it is still ambiguous how this data is generated. Given the potential importance of this process in understanding why the method works, could you provide more clarity on how the negative data is generated and discuss how changing the distribution of negative data affects the results?\n\n- The authors helpfully point out all datasets that were also used as test sets for TabPFN's evaluation since these could potentially be tainted if good practices were not followed on the original TabPFN paper [1].\nI am more concerned however with datasets that were part of the meta-validation set that was used to tune the hyperparameter priors of TabPFN. It seems that one of the datasets used for evaluation of TabEBM (stock) is part of this meta-validation set (Tables 7 and 8 in the appendix of [1]).\nEven though this is only one of 150 datasets that comprise this validation set and it is not clear in the original paper the extent of this finetuning, I would still consider it good practice to avoid using this dataset for evaluating downstream methods for fear of data leakage. I also would have liked the paper to prioritize evaluating on more datasets that weren't used at all in [1]. I understand if [1] might have somewhat exhausted the pool of small datasets in OpenML but evaluating on larger datasets by first subsampling them would be an option. This would also allow for the use of real data as an ideal baseline to measure TabEBM's improvement against.\n\n- For TabEBM and the other methods that learn class conditional densities (CTGAN, TabDDPM, ARF, GOGGLE, TabPFGen) what distribution of classes (i.e., $p(y)$) was used for the synthetic data? This is not specified in the paper so I am assuming it was an empirical estimate of the p(y) from the training data otherwise this would be somewhat unfair when comparing to methods that learn a joint distribution. Furthermore it would be potentially problematic when evaluating a balanced accuracy in a test set if the p(y) of generated data is not the same across methods (for conditional methods).\n\n- How are statistical fidelity results computed? The metrics reported seem to be univariate ones so how are they being used to compare multivariate distributions? Are only marginal distributions being compared? Is it an average over univariate metrics computed for each marginal separately? How is the KS test used on categorical variables? This section of the experiments is sorely lacking in details and the appendices seem to provide no further information. If it is indeed the case that only marginals are being compared I find this evaluation lackluster since it is easy to achieve similar marginal distributions without learning the relationships between different input variables which is the main difficulty in learning multivariate density models.\n\n\n[1] Noah Hollmann, Samuel Müller, Katharina Eggensperger, Frank Hutter. TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second, 2022 (https://arxiv.org/abs/2207.01848)\n\nLIMITATIONS:\nThe authors aknowledge that due to the reliance of the method on a PFN it is currently constrained for use in small datasets.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes TabEBM, a novel data augmentation method designed for low-sample-size tabular classification tasks. TabEBM generates synthetic tabular data using class-specific Energy-Based Models (EBMs) to learn the marginal distribution for each class. Experimental results on various real-world datasets demonstrate that TabEBM improves downstream performance via data augmentation, generates high-fidelity synthetic data, and strikes a competitive balance between accuracy and privacy in data sharing.\n\nSTRENGTHS:\n- Novel approach using class-specific EBMs for tabular data generation.\n- Comprehensive evaluation across multiple datasets, metrics, and downstream tasks.\n- Strong performance, especially in low data regimes.\n- Thorough analysis of statistical fidelity and privacy preservation.\n- Open-source implementation provided.\n\nWEAKNESSES:\n- Scalability Issues: The reliance on TabPFN, which struggles with large sample sizes, limits TabEBM's scalability.\n- Implementation Complexity: The need for class-specific surrogate tasks and the iterative nature of the sampling process may complicate implementation and increase computational overhead.\n\nQUESTIONS:\n- How does TabEBM perform on datasets with  hundreds of features?\n- Could the approach be extended to use other pre-trained models besides TabPFN?\n- How does the method handle highly imbalanced datasets where some classes have significantly fewer samples than others?\n- What is the computational complexity of generating samples compared to other methods?\n- How sensitive is the performance to the choice of hyperparameters in the SGLD sampling process?\n- Clarify the meaning of the word \"fit\" which is used in line 122 of the paper.\n\nLIMITATIONS:\n- Performance may degrade for datasets much larger than those tested.\n- Inherits limitations of the underlying TabPFN model.\n- May not be suitable for datasets with extremely high dimensionality.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors present a new method of tabular data augmentation, called TabEBM. The unique feature of TabEBM is that it creates distinct generative models for each class in a classification problem setting. With extensive and thorough evaluations, the authors prove that TabEBM sets the new state of the art.\n\nSTRENGTHS:\n1. The manuscript is well organized. The text is clear. The quality of the figures is high. \n2. The authors address a very common problem for the broad scientific community. Application of machine learning algorithms on small tabular datasets is limited by design. Introducing TabEBM as an effective data augmentation method, this work can have an extremely broad impact in the future.\n3. The experiments are rich and rigorous. The method is compared against many existing methods on a variety of datasets. The evidence proving its effectiveness is convincing.\n4. The authors claim TabEBM will be available as an open-source library.\n\nWEAKNESSES:\nSome important details on study design, model training, and practical considerations of TabEBM application are missing. See questions.\n\nQUESTIONS:\n1. The definition of a surrogate binary classification task (lines 111-113) might result in class imbalance. How is this potential problem addressed? \n2. Can the authors provide more details on training the models for the surrogate binary classification task?  This could eliminate some other questions.\n3. Have the authors specifically investigated the propensity of TabEBM to generate outliers? How does TabEBM compare to other methods in this regard?\n4. In section 3.2, the authors describe many statistical evaluations performed. However, the multiple hypothesis testing aspect is not described in sufficient detail. How were the p-values adjusted? How strong is the impact of a particular correction method on the final results (tables of C3)?\n5. Can the authors share more insights about the open-source library? This is positioned as one of the key contributions, but the code is not provided and the practical considerations of using the library are not clear from the submission.\n\nLIMITATIONS:\nBriefly discussed in section 4.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces TabEBM, a class-conditional generative method for tabular data augmentation using distinct EBMs for each class. By modeling each class's data distribution individually, TabEBM generates high-quality synthetic data. Experiments demonstrate that using TabEBM for data augmentation improves classification performance across various datasets, particularly small ones.\n\nSTRENGTHS:\nThe paper introduces distinct class-specific EBMs for tabular data augmentation. Moreover, it presents thorough experiments demonstrating effectiveness across various datasets, especially small ones. Additionally, the writing is well-structured and clear.\n\nWEAKNESSES:\n1. On the technical level, TabPFN, EBMs, and SGLD have all been introduced in TabPFGen [48]. The main difference in the proposed TabEBM method is the use of class-specific EBMs. However, this idea is straightforward and lacks novelty.\n2. The most significant difference between tabular tasks and common tasks like images is that they include both continuous and categorical features. However, TabEBM treats all features as continuous, lacking targeted consideration for categorical features.\n3. In TabEBM, although the choice of surrogate binary classifier is arbitrary, not using training-free models like TabPFN will result in training costs being positively correlated with the number of classes, significantly increasing the training overhead.\n4. The datasets used in the experiments have a maximum dimensionality of 77, lacking experimental results on high-dimensional data (e.g., hundreds of dimensions).\n\nQUESTIONS:\n1. Do the datasets used in the experiments only contain continuous features and no categorical features?\n2. The TabEBM method is similar to TabPFGen [48]. Why didn't you use the 18 datasets mentioned in the TabPFGen paper for experiments, but instead chose 6 different datasets (with a smaller test scope) for evaluation?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-11"}