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

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

Paper

References (88)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer XdB46/10 · confidence 4/52024-07-06

Summary

The 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. In order to avoid training these classifiers on the data itself, the authors take advantage of a pre-trained Prior-Data Fitted Network (PFN). The 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.

Strengths

- The paper is well-structured and clearly written, effectively communicating the key ideas and concepts. - 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. - 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. - 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.

Weaknesses

- 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$. My 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. - 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. Personally, 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. [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)

Questions

- 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. - 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? - The details of how the negative examples are generated are not clear. The text simply states: "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? - 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]. I 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]). Even 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. - 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). - 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. [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)

Rating

6

Confidence

4

Soundness

2

Presentation

3

Contribution

3

Limitations

The authors aknowledge that due to the reliance of the method on a PFN it is currently constrained for use in small datasets.

Reviewer iojh5/10 · confidence 4/52024-07-09

Summary

The 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.

Strengths

- Novel approach using class-specific EBMs for tabular data generation. - Comprehensive evaluation across multiple datasets, metrics, and downstream tasks. - Strong performance, especially in low data regimes. - Thorough analysis of statistical fidelity and privacy preservation. - Open-source implementation provided.

Weaknesses

- Scalability Issues: The reliance on TabPFN, which struggles with large sample sizes, limits TabEBM's scalability. - Implementation Complexity: The need for class-specific surrogate tasks and the iterative nature of the sampling process may complicate implementation and increase computational overhead.

Questions

- How does TabEBM perform on datasets with hundreds of features? - Could the approach be extended to use other pre-trained models besides TabPFN? - How does the method handle highly imbalanced datasets where some classes have significantly fewer samples than others? - What is the computational complexity of generating samples compared to other methods? - How sensitive is the performance to the choice of hyperparameters in the SGLD sampling process? - Clarify the meaning of the word "fit" which is used in line 122 of the paper.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

- Performance may degrade for datasets much larger than those tested. - Inherits limitations of the underlying TabPFN model. - May not be suitable for datasets with extremely high dimensionality.

Authorsrebuttal2024-08-12

Additional clarifications and results on imbalanced datasets and large datasets

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. | Datasets `(#Minority:#Majority)` | Baseline | SMOTE | TVAE | CTGAN | TabDDPM | TabPFGen | TabEBM (Ours) | | --- | --- | --- | --- | --- | --- | --- | --- | | **steel** | | | | | | | | | `(50:50)` | 88.79 | 77.74 | 74.21 | 77.61 | 83.78 | 89.67 | **93.77** | | `(20:80)` | 74.33 | 66.56 | 68.10 | 61.52 | 69.06 | 83.38 | **84.65** | | `(10:90)` | 60.74 | 60.76 | 54.85 | 53.50 | 58.46 | 70.36 | **73.02** | | `(5:95)` | 52.29 | 56.11 | 51.12 | 51.30 | 51.17 | 64.84 | **65.95** | | **stock** | | | | | | | | | `(50:50)` | 88.64 | **90.21** | 88.67 | 85.52 | 89.23 | 88.16 | 90.16 | | `(20:80)` | 85.34 | 88.95 | 85.87 | 85.37 | 87.16 | 88.72 | **89.38** | | `(10:90)` | 76.05 | 83.86 | 74.60 | 76.93 | 78.81 | 83.52 | **84.92** | | `(5:95)` | 65.98 | 75.24 | 62.66 | 61.11 | 68.69 | 75.56 | **77.68** | 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. | N_real | Baseline | SMOTE | TVAE | CTGAN | TabDDPM | TabPFGen | TabEBM (Ours) | Improvement vs Baseline by TabEBM (%) | | --- | --- | --- | --- | --- | --- | --- | --- | --- | | 50 | 72.40 | 76.40 | 55.332 | 54.80 | 62.94 | N/A | **78.90** | +6.50 | | 100 | 82.42 | 84.35 | 66.00 | 69.49 | 76.34 | N/A | **86.01** | +3.59 | | 200 | 87.54 | 89.29 | 78.37 | 82.44 | 82.53 | N/A | **89.77** | +2.23 | | 500 | 92.96 | 93.69 | 90.09 | 91.48 | 91.24 | N/A | **93.76** | +0.80 | | 1000 | **96.37** | 96.21 | 93.61 | 95.36 | 94.56 | N/A | 96.30 | -0.07 | | 2000 | 97.76 | 96.84 | 96.62 | 97.10 | 97.13 | N/A | **97.83** | +0.07 | | 3000 | 98.20 | 98.28 | 97.60 | 97.60 | 97.73 | N/A | **98.35** | +0.15 | | 4000 | 98.51 | **98.59** | 98.11 | 98.00 | 98.46 | N/A | 98.55 | +0.04 |

Reviewer Cc3y7/10 · confidence 4/52024-07-09

Summary

The 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.

Strengths

1. The manuscript is well organized. The text is clear. The quality of the figures is high. 2. 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. 3. 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. 4. The authors claim TabEBM will be available as an open-source library.

Weaknesses

Some important details on study design, model training, and practical considerations of TabEBM application are missing. See questions.

Questions

1. The definition of a surrogate binary classification task (lines 111-113) might result in class imbalance. How is this potential problem addressed? 2. Can the authors provide more details on training the models for the surrogate binary classification task? This could eliminate some other questions. 3. Have the authors specifically investigated the propensity of TabEBM to generate outliers? How does TabEBM compare to other methods in this regard? 4. 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)? 5. 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.

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

Briefly discussed in section 4.

Reviewer tdoQ5/10 · confidence 4/52024-07-11

Summary

The 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.

Strengths

The 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.

Weaknesses

1. 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. 2. 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. 3. 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. 4. The datasets used in the experiments have a maximum dimensionality of 77, lacking experimental results on high-dimensional data (e.g., hundreds of dimensions).

Questions

1. Do the datasets used in the experiments only contain continuous features and no categorical features? 2. 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?

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

N/A

Authorsrebuttal2024-08-12

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

Reviewer tdoQ2024-08-13

I appreciate the additional experiments conducted by the author. These evaluations addressed my concerns regarding the experiment scale and performance for datasets with many categorical features. I have increased my overall rating for this paper from Reject to Borderline Accept.

Reviewer XdB42024-08-07

Response

I thank the authors for the clarifications and for taking the time to run these additional experiments. I must confess that the reason this approach works is still somewhat mysterious to me and I would like to have a better theoretical justification for it. However, the linked blogpost and the author's additional experiments do shed some light on the matter and seem to suggest that this may be due to some property of TabPFN. The authors have, at least, successfully convinced me that their surrogate task procedure works and is somewhat reliable and I believe these are valuable empirical results. I believe the clarification on the generation of negative data should be in the paper. Since I consider my concerns mostly addressed, I am happy to raise my score to a 6.

Reviewer Cc3y2024-08-11

Thanks to the authors for the clarifications and additional experiments! I believe my initial evaluation was fair, so my score remains the same.

Reviewer iojh2024-08-11

Thank you to the authors for clarifying and conducting additional experiments. These efforts have significantly enhanced the quality of the research. However, I believe my initial evaluation remains valid, so my score remains unchanged.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC