TabDiff: a Mixed-type Diffusion Model for Tabular Data Generation

Authors: Juntong Shi, Minkai Xu, Harper Hua, Hengrui Zhang, Stefano Ermon, Jure Leskovec
Primary area: applications to computer vision, audio, language, and other modalities
Keywords: Tabular Representative Learning, Generative Models, Diffusion Models

Abstract:
Synthesizing high-quality tabular data is an important topic in many data science tasks, ranging from dataset augmentation to privacy protection. However, developing expressive generative models for tabular data is challenging due to its inherent heterogeneous data types, complex inter-correlations, and intricate column-wise distributions. In this paper, we introduce TabDiff, a joint diffusion framework that models all mixed-type distributions of tabular data in one model. Our key innovation is the development of a joint continuous-time diffusion process for numerical and categorical data, where we propose feature-wise learnable diffusion processes to counter the high disparity of different feature distributions. TabDiff is parameterized by a transformer handling different input types, and the entire framework can be efficiently optimized in an end-to-end fashion. We further introduce a mixed-type stochastic sampler to automatically correct the accumulated decoding error during sampling, and propose classifier-free guidance for conditional missing column value imputation. Comprehensive experiments on seven datasets demonstrate that TabDiff achieves superior average performance over existing competitive baselines across all eight metrics, with up to $22.5\%$ improvement over the state-of-the-art model on pair-wise column correlation estimations. Code is available at https://github.com/MinkaiXu/TabDiff.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730335248716

thread_type:
official_review

<thread_object>
creator: reviewer_buu2

<content>
summary:
This paper introduces TABDIFF, a joint diffusion framework which aims to handle data with both discrete and continuous data types. The authors validated the effectiveness of this framework through comprehensive experiments.

strengths:
1.The method introduces a new approach by jointly learning the diffusion processes for discrete and continuous features, which is not explored in the previous work.

2. The adaptively learnable noise schedules look novel.

3.The paper provides extensive experimental evaluations across different benchmark datasets, thoroughly validating the method’s performance.

weaknesses:
In the title of this paper, it is claimed that this paper addresses multi-modal distributions. However, the paper does not clearly specify what multi-modal distribution it handles. Does it only refer to distributions with both discrete and continuous variables?

questions:
NA
</content>
</thread_object>

</thread>
<thread>
creation date:
1730475611270

thread_type:
official_review

<thread_object>
creator: reviewer_kzhm

<content>
summary:
The paper introduces TabDiff, an innovative multi-modal diffusion model designed to generate high-quality synthetic tabular data. It addresses the challenges of heterogeneous tabular data, which includes both numerical and categorical features with complex interdependencies. TabDiff employs a joint continuous-time diffusion process with feature-wise learnable components, allowing it to handle different data modalities while adapting to the varying distributions of individual features.

strengths:
1. Tabdiff introduces a unified joint continuous-time diffusion process with feature-wise learnable components for handling both numerical and categorical data is innovative and allows the model to adapt to the specific characteristics of different data distributions.
2. The paper provides a thorough evaluation of TabDiff on a wide range of metrics, offering a deep and detailed assessment of the model performance from multiple dimensions, which is highly commendable.
3. The paper emphasizes the importance of TabDiff in generating high-quality synthetic data that enables effective model training while ensuring privacy preservation. The authors validate this effectiveness through the use of various metrics, demonstrating that the model not only maintains data fidelity but also safeguards against privacy risks. It's a practical trial for addressing the growing need for privacy-conscious data generation in sensitive domains.

weaknesses:
1. Although the paper does not explicitly highlight it, the complexity of TabDiff, particularly in handling high-dimensional features through diffusion, could result in increased computational costs compared to simpler models. It would be good to provide a comparison of generation or training speed metrics with other baselines to better illustrate this trade-off.

2. The performance drop of TabDDPM on the News and Diabetes datasets is surprising, especially given that other diffusion-based methods maintain stable results. This discrepancy is beyond expected boundaries, with precision and recall even dropping to 0 in Table 7 and Table 8, and similarly poor performance across other tables for these datasets. 
Thus, while I appreciate the authors' comprehensive comparison across various metrics with baselines, the paper lacks sufficient analysis and discussion of the experimental results, particularly in explaining the performance gaps between TabDiff and the baselines. It would strengthen the paper if the authors provided more detailed discussion, and offered insights into the experimental setup of the baselines.

3. The paper lacks sufficient explanation of the experimental setup for these baselines, which makes it harder to interpret the results. Providing more clarity on these aspects would strengthen the case for TabDiff's performance and offer a more comprehensive understanding of the comparative results.

questions:
See weaknesses, I will raise the score if all the concerns are addressed.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730522571324

thread_type:
official_review

<thread_object>
creator: reviewer_ud32

<content>
summary:
This paper introduces an approach called TABDIFF, which develops a joint continuous-time diffusion process for mixed-type data in tabular datasets. The authors address the challenge of high disparity between different feature distributions by proposing feature-wise learnable diffusion processes. The key contribution is a hybrid diffusion process to handle numerical and categorical features.

strengths:
**1** The experiments are extensive, demonstrating the efficiency of the proposed method.

**2** Hybrid diffusion models are employed to handle mixed-type data in tabular datasets.

weaknesses:
**1** The authors claim that they explicitly tackle the feature-wise heterogeneity issue in the multi-modal diffusion process. However, the details about how to handle this problem is missing.

**2** Even though experiments show the efficiency of the proposed method, it lacks of the motivation of considering a masked diffusion for categorical features.

questions:
**Q1** The title of this paper is somewhat confusing. Before reading it, I assumed it referred to a diffusion model capable of generating various types of data structures, such as tabular data and images. After reading, I believe the term "modality" may refer to either of two scenarios: (1) multiple modalities within a specific variable (a column in a tabular dataset); (2) columns of a tabular dataset that are mixed-type. If the latter is the intended meaning, the paper should demonstrate how the proposed method differs from TabDDPM, as TabDDPM also handles mixed-type data generation in tabular datasets. It is inappropriate to use such an ambitious title.


**Q2** The proposed method utilizes the masked diffusion for categorical features. What is the point of using masked diffusion?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730548779897

thread_type:
official_review

<thread_object>
creator: reviewer_ztt5

<content>
summary:
This work proposed a multi-modal diffusion framework for generating high-quality synthetic data. A key innovation is the development of a joint continuous-time diffusion process for numerical and categorical data, where authors propose learnable column-wise noise schedules to counter the high disparity of different feature distributions. Extensive experiments are conducted to demonstrate the superior capacity of proposed approach.

strengths:
This manuscript demonstrates advantages in multiple aspects. It provides a clear and comprehensive introduction of the research problem, and conducts extensive experiments on multiple data sets and evaluation criteria.

Additionally, the paper elaborates on relevant algorithms and techniques. The algorithmic sections present some innovative and original findings, contributing to advancements in the field of machine learning and offering new insights to the academic community.

weaknesses:
1. Some necessary explanations are lacking in the description of algorithm provided in the Method and Appendix, requiring more detailed descriptions to enhance the readability of the paper.

2. The description of categorical column needs an example to illustrate the method more clearly.

3. Formula (10) is missing a parenthesis.

4. Why are formula (10) and formula (11) constructed this way? Is there a source? If not, can you explain the construction idea?

questions:
See weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730680416821

thread_type:
official_review

<thread_object>
creator: reviewer_gkxc

<content>
summary:
While extensive work has attempted to solve the task of tabular data generation, they often face different problems, such as the failure of cross-modal modeling and the additional cost of training auto-encoders. To this end, this paper proposes TabDiff, a diffusion-based model that unifies modal heterogeneity within one single approach. The proposed method, built upon several existing techniques, models numerical and categorical attributes separately and introduces a learnable noise scheduler that further facilitates the sampling process. Extensive experiments verify the effectiveness of TabDiff. Ablation studies also provide additional insights.

strengths:
1. The paper is well-written and well-structured. It is self-content and easy to follow.
2. The problem investigated in this paper is well-motivated and practical.
3. While most of the techniques have been investigated in prior work, this work seamlessly integrates them and adapts to the tabular data generation task.
4. Extensive experiments sufficiently verify the effectiveness of the proposed approach, and ablation studies further outline the design choices and the benefits of the learnable noise schedulers.

weaknesses:
1. Despite the claim of cross-modal modeling, the proposed method currently still models numerical and categorical values separately, which might make it difficult to synthesize more complicated tabular datasets. 
2. While I generally appreciate the effort and the performance, I am not fully convinced by the claim of the multi-modal diffusion framework. The experiments in Tables 1, 2, and 3 all suggest comparable performance compared to TabSyn. As TabSyn has been published in ICLR 2024, the authors should clearly state the main difference and closely compare the two works.
3. Lack of motivation for schedulers. What is the motivation for choosing power-mean and log-linear schedulers? The choice looks random to me.

questions:
1. Do numerical and categorical value modeling help each other, or are they just modeled separately?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732529621273

thread_type:
public_comment

<thread_object>
creator: ~markus_mueller1

<content>
comment:
Dear authors,

Thank you for this very interesting read and your contribution to diffusion models for tabular data. It is great to see progress on noise schedules specifically designed for tabular data.

We have a specific interest in your paper as it is closely related to our work on diffusion models for tabular data [1]. In our paper, we cast both categorical and continuous features into continuous space, to be able to apply Gaussian noise with adaptive noise schedules to both feature types.  

Below we list a couple of questions that would help us (and hopefully you as well) to better understand the impact of some modeling choices:

- We found that our modeling choices help to balance model capacity allocation across features. Did you also find a similar balancing impact of using feature-specific noise schedules in your work? Would the balancing across features be sensitive to how different the scales of numerical and categorical losses are?
- In your work, the categorical features are generated via unmasking and cannot change afterwards. In that case, have you seen any deterioration in results if categorical features depend a lot on the fine-grained details of the continuous features (which may be generated only after a categorical feature has already been unmasked)?
- Your choice of letting $\lambda_{\text{num}}$ linearly decay as training proceeds is interesting. Essentially this seems to imply that towards the end of training parameter updates will be devoted towards modeling categorical features. In our own experiments, however, we found that modeling categorical features is rather easy compared to getting the fine details of continuous features right. What is the rationale behind your choice?

---

[1] Mueller, et al. (2024) "Continuous Diffusion for Mixed-Type Tabular Data".
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper explores “multi-modal” diffusion generative models for tabular data, where “multi-modal” refers to numerical and categorical tabular data only. The work develops a joint continuous-time process with learned columnwise noise schedules that help to model the heterogeneous features typical in tabular data. 

Reviewers initially had a range of concerns, but these were generally addressed during the discussion.

While I am recommending acceptance as a poster, I have two requirements that the authors must address in their camera-ready version:

1) Change the title and writing throughout the paper to deemphasize “multi-modal”, as this term is not being used in a standard sense. The paper only considers the tabular modality, which can have mixed-type columns (numerical and categorical). Hence changing “multi-modal data” to “mixed-type tabular data” for example would better communicate the scope of the work.

2) Reduce the length of the main text to fit within the ICLR page limit.

Authors, please post a comment on this thread once you have made these changes and I will review. If I do not see your post and changes by the camera-ready deadline I will raise the issue to SAC/PC. Thank you for cooperating!