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

Thank you for the detailed review and constructive feedback! We address all your questions and comments below, and provide a rebuttal PDF in the general rebuttal. Due to space limits, we summarise the new results. We will update our manuscript with additional clarifications and results.

>Q3a: Clarify the generation of the negative samples

The surrogate classification task is to determine if a sample belongs to class $c$ by comparing $\mathcal{X}_c$ against a set of "negative samples" $\mathcal{X}_c^{\text{neg}}$. We label the true samples $\mathcal{X}_c$ as 1 and the negative samples $\mathcal{X}_c^{\text{neg}}$ as 0, and train TabPFN on this datasets, resulting in a class-specific binary classifier used to define the $E_c(x)$ (using Equation 3)

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 $\alpha\sigma_d$ or $-\alpha\sigma_d$, where $\alpha$ is a fixed constant and $\sigma_d$ is the standard deviation of dimension $d$. For example, in $R^3$, a negative sample might have coordinates $[\alpha\sigma_1, \alpha\sigma_2, -\alpha\sigma_3]$. In the paper we use four negative samples with $\alpha = 5$, placing them far from the real data.

>Q3b: How does the distribution of negative data affect the results?

We found it essential to use negative samples far from the real data, to allow the binary classifier to easily differentiate between $\mathcal{X}_c$ and $\mathcal{X}_c^{\text{neg}}$.

Figure F1 from the PDF shows a new experiment varying the distribution 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, affecting its logits magnitude and making them unsuitable for density estimation (we further investigate the logits in our next answer to Q1/W1).

TabEBM is robust to the distribution of negative samples if they are easily distinguishable from real data. We ran two new experiments, keeping the negative samples at the corners of a hypercube and varying their distance and number. When varying the per-dimension distance of those samples $\alpha \in [0.1, 5]$, TabEBM provides consistent improvements between 2.85-3.18%. Additionally, our results in Table R2 (in our answer to Q1 from reviewer Cc3y) demonstrate that TabEBM performs similarly regardless of the number of negative samples in the surrogate binary tasks.

>Q1/W1: Why can TabPFN's logits be useful for energy estimation?

Indeed, we agree with your observation! As the TabPFN classifier is trained on different class-specific surrogate tasks, it learns different logits for each task. We found it essential to place the negative samples far from the real samples, because TabPFN's confidence depends on the distance to the training data [1], as it was pre-trained to approximate Bayesian inference.

We present new experiments in Figure F2 from the PDF. The top row shows that as the distance to the real data increases, the logit $f(x)[1]$ for the real data smoothly decreases until the two logits become similar. Thus, the classifier predicts is uncertain in these far-away regions. As maximum logits decrease, TabEBM's inferred density drops significantly (as shown on the bottom row) because $p_c(x) \propto (\exp(f(x)[0]) + \exp(f(x)[1]))$. One could possibly fine-tune TabPFN on the surrogate tasks to further improve the logits for density estimation.

>Q2: Why fit p(x) rather than p(x|y=1) in the surrogate binary tasks?

Both approaches lead to virtually identical results due to the design of our surrogate task. The new results in Figure F2 (bottom row) from the PDF shows that the TabEBM's inferred energy using $p(x)$ (in red) is essentially identical to the energy of $p(x|y=1)$ (in grey), especially near the real data. As the SGLD sampling starts near the original points (line 676), it will visit the neighbourhood of the real samples, where the energy surfaces are virtually identical, leading to similar results.

>Q4: Evaluating more datasets

We run new experiments on six leakage-free datasets from UCI, with 1,000-9105 samples of 7-42 features. Table R2 (in our response to Q2 for tdoQ) shows that TabEBM consistently outperforms the baseline and all other benchmark methods.

>Q5: What's the class distribution of the synthetic data?

The synthetic data has the same class distribution as the real training data.

>Q6: Is statistical fidelity computed over the univariate marginals?

Yes, these metrics are univariate, and we computed them using the open-source library Synthcity. We acknowledge the reviewer's point about the imperfection of univariate metrics, although evaluating generators' ability to capture the joint feature relationships remains underexplored [2].

>W2: How does tuning the generators/predictors affect the ranking?

We run new experiments tuning three generators on three datasets using the average validation accuracy of six downstream classifiers. The ranges are: SMOTE ($k \in \{3, 5, 10\}$), TVAE ($\text{lr} \in \{5e-4, 1e-3, 5e-3\}$) and CTGAN ($\text{lr} \in \{1e-3, 5e-3, 1e-2\}$). TabEBM remains the most competitive data augmentation method, improving accuracy by +2.25, followed by SMOTE (+1.81) and CTGAN (-5.71). We can provide the full results in the discussion.

We also tuned the downstream predictors, and the new results in Table R7 from the PDF show that TabEBM remains the best-performing method, providing the largest improvements for data augmentation, even after tuning the predictors.

Thank you again for your constructive review! We would appreciate it if you would consider raising your score in light of our response.

References:

• [1] McCarter, What exactly has TabPFN learned to do?, ICLR Blogposts, 2024
• [2] Tu, Ruibo, et al., 2024, (https://arxiv.org/abs/2406.08311)

Thank you for your thoughtful review! We address all your questions and comments below. Due to limited space, we summarised the new results. We will update the manuscript to include the complete new experiments and clarifications.

>Q3: How does TabEBM perform on highly imbalanced datasets?

We conducted new experiments, adjusting the class imbalance on two binary datasets (with $N_\text{real}=100$). We keep the setup from Section 3 and report the balanced accuracy averaged over six downstream predictors. Table R4 shows that under high class imbalanced, TabEBM outperforms the Baseline and SMOTE, which is a method specifically designed to handle imbalanced datasets.

Table R4. Test balanced classification accuracy (%) varying the class imbalance.

>Q5: How sensitive is TabEBM to the hyperparameters of the SGLD sampling?

We run new experiments varying two key hyperparameters of SGLD on the “biodeg” dataset with $N_{\text{real}}=100$. Specifically, we vary the step size $\alpha_{\text{step}}$ and the noise scale $\alpha_{\text{noise}}$, reporting the accuracy averaged over six downstream predictors.

Table R5 shows that TabEBM is stable to these hyperparameters, as the difference between the highest and lowest accuracy is less than 1.5%. Note that increasing $\alpha_{\text{noise}}$ (which is added at each SGLD step) is expected to degrade performance because we standardized the data to have unit standard deviation.

Table R5. Classification accuracy (%) of TabEBM with different SGLD settings

>Q6: Clarify the meaning of the word "fit" on line 122

We use the term “fit” to mean “train” the in-context model TabPFN. For TabPFN, “training” is similar to the K Nearest Neighbour algorithm, where "training" means defining the training dataset used for inference. Thus, we simply update TabPFN's training datasets and the model in inference mode, thus never updating its parameters.

>Q4: How does TabEBM's computational complexity compare to other methods?

Figure 4 from the paper illustrates the trade-off between accuracy and the time required for training and generating 500 synthetic samples. The results show that TabEBM is the fastest method for generating data besides SMOTE. Other methods are 3-30 times slower than TabEBM. Also, TabEBM surpasses all other methods in improving downstream accuracy through data augmentation.

>W2: Class-specific surrogate tasks and iterative sampling can complicate implementation and increase computational overhead

On implementation complexity: TabEBM simplifies implementation by eliminating the need for dedicated training, pipelines, protocols, and resources—making it ready to use for generating data. We provide an open-source implementation of TabEBM, allowing users to generate data with just two lines of code for new datasets. We believe this ease of use will encourage the widespread application of tabular data augmentation on small datasets.

On the computational overhead: As mentioned in our answer to Q4 and Figure 4, TabEBM is computationally efficient, being the fastest method for generating data besides SMOTE —while outperforming all competing methods.

>Q2: Can TabEBM use other pre-trained models besides TabPFN?

Yes, TabEBM can be used with other pre-trained models besides TabPFN. In the paper, we use TabPFN because it is the only tabular model of this type with open-sourced weights. However, TabEBM is a general method for transforming classifiers into class-specific generators and can use any gradient-based in-context classifier that computes logits (through Equation 3). As new tabular foundational models are developed [1], they can be readily integrated into TabEBM.

>Q1/W1/L1/L2/L3: On the scalability and limitations of TabEBM

TabEBM is a method that uses an underlying binary classifier, thus its capabilities and limitations are directly influenced by the in-context model it employs (TabPFN in our case), as discussed in the paper (Lines 317-323). Since TabPFN handles only up to 100 features, this limitation is also inherited by TabEBM.

Our new experiment (Table R6) shows that TabEBM can be applied to larger datasets beyond TabPFN's limitation of 1000 samples. On larger datasets TabEBM still outperforms other generators (not shown due to space limits), but training on real data alone appears sufficient. This highlights TabEBM's usefulness in fields with limited training samples. Note that despite TabPFN's 10-class limit, TabEBM can handle unlimited classes by using TabPFN only for surrogate binary tasks.

As foundational models for tabular data evolve [1], new models that can accommodate more features are anticipated. Integrating these models into the TabEBM framework will enable it to handle high-dimensional datasets, thus increasing its versatility and utility.

Table R6. Classification accuracy (%) comparing data augmentation with increased real data availability on the “texture” dataset

Thank you again for your thoughtful review! We would appreciate it if you would consider raising your score in light of our response.

References:

• [1] Boris van Breugel, Mihaela van der Schaar, Why Tabular Foundation Models Should Be a Research Priority, ICML 2024 (https://arxiv.org/abs/2405.01147)

Dear Reviewer iojh,

Thank you for acknowledging the additional experiments and clarifications we provided. We are glad that you found that our efforts have significantly enhanced the quality of our research. We are also pleased that, in your original review, you highlighted several strengths of our work, including its novelty, comprehensive evaluation, strong performance, and open-source implementation.

We conducted additional experiments to further support our argument, which we believe will interest the reviewer. Specifically, we evaluated four additional generative models—TVAE, CTGAN, TabDDPM, and TabPFGen—addressing your questions on handling imbalanced datasets and the effectiveness of data augmentation on large datasets.

• Table R5-extended demonstrates that TabEBM outperforms these methods, especially on highly imbalanced datasets. The results indicate that TabEBM is particularly effective for class balancing, and we will highlight this application in the revised manuscript.
• Table R6-extended suggests that data augmentation may not be necessary when ample real data is available for model training. As the availability of real data increases, training downstream predictors solely on this real data can guarantee optimal performance. Existing methods, such as TVAE, CTGAN and TabDDPM, tend to decrease performance on smaller datasets, where performance enhancement is most needed, while TabEBM consistently delivers the largest improvements. This suggests that data augmentation is beneficial primarily in data-scarce scenarios, where TabEBM excels. We will include these results in the updated manuscript.

Given the new results and your positive feedback, we kindly ask if you might reconsider your initial score in light of these results. If there are any remaining concerns, please let us know, as we would be more than happy to discuss them further while the discussion period is open. We would greatly appreciate your feedback in order to further improve our paper.

Thank you again for your time and consideration,

Authors

Table R5-extended. Test balanced classification accuracy (%) varying the class imbalance. We aggregated the performance over six downstream predictors. TabEBM is effective for datasets with high class imbalance.

Table R6-extended. Test balanced classification accuracy (%) aggregated over six downstream predictors, comparing data augmentation with increased real data availability of the “texture” dataset. TabPFGen is not applicable because it supports only datasets with up to 10 classes.

Thank you for your thoughtful review! We address all of your questions and comments below. Due to space constraints, we summarised the new results. We will update the manuscript to include the complete new experiments and clarifications.

>Q2: Provide details on training models for the proposed surrogate binary classification tasks

In TabEBM, for each class $c$, we train a class-specific EBM, $E_c(x)$, using exclusively the data from that class, denoted as $\mathcal{X}_c$. The energy function $E_c(x)$ is derived by adapting the logits from a binary classifier, as outlined in Equation 3. To train this classifier with only class-specific data, we create surrogate binary classification tasks.

Our surrogate classification tasks determine if a sample belongs to class $c$ by comparing $\mathcal{X}_c$ against a set of "negative samples," $\mathcal{X}_c^{\text{neg}}$. We intentionally create these negative samples at a distance of five standard deviations from the original data to ensure the binary classifier can easily differentiate between the real class data and these artificially distant samples.

We label the true class samples $\mathcal{X}_c$ as 1 and the negative samples $\mathcal{X}_c^{\text{neg}}$ as 0, forming a combined dataset $\mathcal{D}_c = (\mathcal{X}_c \cup \mathcal{X}_c^{\text{neg}}, \{1\}^{|\mathcal{X}_c|} \cup \{0\}^{|\mathcal{X}_c^{\text{neg}}|})$. Following this, we train the TabPFN classifier on $\mathcal{D}_c$, resulting in a binary classifier specific to class $c$. This classifier is then used to define the energy function $E_c(x)$ using Equation 3, from which we can generate data using SGLD.

>Q1: Does the surrogate binary classification task result in class imbalance issues?

No, the imbalance between the "negative samples" and the real samples used in the surrogate binary classification task has a negligible impact on TabEBM's performance.

We conducted a new experiment to evaluate the effects of varying the ratio $\|\mathcal{X}^{\text{neg}}_c\|:\|\mathcal{X}_c\|$. To ensure reliable outcomes, we maintained a consistent ratio across all classes, keeping the same proportion of negative samples for each class. We varied the number of negative samples $\|\mathcal{X}^{\text{neg}}_c\|$ to represent various ratios of $\|\mathcal{X}_c\|$, thus simulating both balanced and highly imbalanced scenarios.

The results, presented in Table R3, computed across six datasets with $N_{\text{real}} = 100$ real samples, demonstrate that TabEBM is robust to potential imbalances in the surrogate binary tasks.

Table R3. Classification accuracy (%) on data augmentation, showing the impact of the number of negative samples $\|\mathcal{X}^{\text{neg}}_c\|$ in TabEBM across six datasets. TabEBM's performance is robust regarding the number of negative samples.

>Q3: TabEBM's propensity to generate outliers

We provide fair and thorough comparisons by (i) including a wide range of evaluation metrics and (ii) utilising a large synthetic set (Lines 183-184). Given our extensive test scope (i.e., eight datasets $\times$ five sample sizes), it is computationally infeasible to employ complex and specialised methods to detect outliers in synthetic data. However, the downstream accuracy and statistical fidelity metrics are highly indicative of the distribution shifts between real and synthetic data. Notably, Figure 3 illustrates that TabEBM consistently outperforms the baseline across various sample sizes and class numbers. We believe these two metrics adequately demonstrate that TabEBM is more stable to generate in-distribution samples than benchmark methods.

>Q4: Details on the multiple hypothesis testing in Section 3.2

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 statistical fidelity metrics with the open-source implementations from the well-established open-source library Synthcity. However, the previous studies 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. We will further clarify this in the revised manuscript.

>Q5: The code for the open-source library is not provided

We included the TabEBM library as a zip file in the Supplementary material. The library has two core functionalities:

1. Generate synthetic data: The library can generate data that can be used as additional training material for data augmentation. We have included a demo notebook, TabEBM_generate_data.ipynb, in the attached codebase. This notebook demonstrates how to use the generated data for data augmentation and shows a toy example comparing the quality and position of the synthetic data with the real data.

from TabEBM import TabEBM

tabebm = TabEBM()
augmented_data = tabebm.generate(X_train, y_train, num_samples=100)
# augmented_data[class_id] = generated data for a specific `class_id`

2. Compute & visualise the energy function: We allow users to compute TabEBM's energy function and the unnormalised data density. The demo notebook, TabEBM_approximated_density.ipynb, shows TabEBM's energy under conditions of data noise and class imbalance.

We will make the TabEBM library freely available after publication. It is easy to use, domain-agnostic, and requires no training, making it suitable for data augmentation, especially for small-sized datasets.

Thank you again for your thoughtful review! We would appreciate it if you would consider raising your score in light of our response.

Thank you for your thoughtful review! We address all of your questions and comments below. Due to space constraints, we summarised the new results. We will update the manuscript to include the complete new experiments and clarifications.

>Q1/W2: Do the datasets contain categorical features? How does TabEBM handle them?

Yes, we evaluated TabEBM on datasets containing both continuous and categorical features. Specifically, Table R1 shows the two considered datasets with mixed feature types. In lines 639-640, we also clarify that TabEBM encodes the categorical features with leave-one-out target statistics. As the reviewer mentioned, the evaluation results (Table 1) show that TabEBM is effective across various datasets, including those containing categorical features.

Table R1. Number of categorical features in the datasets considered.

>Q2: Why not evaluate the 18 datasets in the TabPFGen paper?

Firstly, we wanted to avoid data leakage and ensure a robust comparison by also evaluating datasets not in TabPFN's meta-test (Appendix F.3 in its original paper). As the 18 datasets in the TabPFGen paper were all used for constructing and evaluating TabPFN, we also use different, leakage-free datasets. Secondly, our test scope is not smaller than that of the TabPFGen paper. We evaluate a wide range of test scopes, considering different datasets and varying the availability of real data. We also include datasets with categorical data, unlike TabPFGen, which investigates only numerical features. Our setup (i.e., eight datasets × five sample sizes) leads up to 33 different test cases. For reference, ARF has the most test cases among prior studies, with 20 test cases, while TabPFGen follows with 18 test cases. Thirdly, we further provide new results on six more leakage-free datasets for UCI with 7-42 features, some including categorical features. We set $N_\text{real}=100$ and provide results for the Top-5 benchmark methods. Table R2 shows that TabEBM continues to outperform all other benchmark methods.

Table R2. Classification accuracy (%) comparing data augmentation on six new leakage-free datasets.

>W1: The core ideas “have all been introduced in TabPFGen” … "this idea is straightforward and lacks novelty”

TabEBM makes a novel contribution to the field of tabular data augmentation by being the first to introduce class-specific generators, as we discussed in our related work, as well as acknowledged by all other reviewers.

Specifically, TabEBM is the first method to create distinct class-specific EBMs. It consists of individual models—one for each class—designed to learn distinct marginal distributions for the inputs associated with each class. TabEBM's class-specific generation is not available to TabPFGen, and it is only possible due to our proposed surrogate binary tasks, which enable the creation of a binary classification task for each class and the obtaining of the class-specific EBM via Equation 3. The results demonstrate that training class-specific EBMs significantly improves the inferred energy over TabPFGen under high levels of noise (Figure 7) and high data imbalance (Figure 8). Unlike TabPFGen, which supports only 10 classes, TabEBM can handle an unlimited number of classes. As you rightly point out, our extensive experiments demonstrate that our method, with its class-specific generation, outperforms all other methods in accuracy, statistical fidelity, and privacy.

In addition to the method, our paper makes two additional key contributions. First, we conduct the first extensive analysis of tabular data augmentation across various dataset sizes. Our benchmark reveals that existing tabular augmentation methods may improve performance only on small datasets, a previously uninvestigated issue. Second, we will release our method as an open-source library for tabular data augmentation, allowing users to generate data immediately without additional training.

>W3: “In TabEBM … not using training-free models like TabPFN will result in … significantly increasing the training overhead.”

Because we tackle classification problems on small-size datasets, training models from scratch is not practical, often leading to suboptimal performance due to overfitting. Instead, using pre-trained models, such as TabPFN, is integral to our approach. From a practical standpoint, using pre-trained models makes TabEBM immediately usable. A significant limitation of existing tabular generators is their need for training, which can be time-consuming, error-prone, and typically requires GPUs. In contrast, our TabEBM library (included with the submission) allows users to start generating data immediately without additional training or tuning.

>W4: Lacking experimental results on high-dimensional data

As noted in Lines 317-323, scaling up in-context tabular models for high-dimensional data remains an unresolved challenge and is not the primary focus of our study. Instead, TabEBM is designed to demonstrate the effectiveness of using pretrained in-context tabular classifiers for generating tabular data. As acknowledged by the reviewer, and further shown in our response to Q2, our "thorough experiments" demonstrate TabEBM's effectiveness.

Thank you again for your thoughtful review! We would appreciate it if you would consider raising your score in light of our response.

Dear Reviewer tdoQ,

Thank you once more for taking the time to provide your feedback on our work. In light of our response and the new experiments, and with the discussion window closing soon, please let us know if you have any further questions or if there is anything else we can clarify. If not, we would appreciate it if you would consider updating your review based on our rebuttal.

Thank you,

Authors