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

TL;DR

TabDiff is a novel diffusion-based generative model designed for high-quality tabular data synthesis, effectively handling mixed data types and complex inter-column relationships through a joint continuous-time diffusion process and transformer architecture.

Contribution

It introduces a unified diffusion framework for mixed-type tabular data, including feature-wise learnable processes and a mixed-type sampler for improved data generation quality.

Findings

Achieves up to 22.5% improvement over state-of-the-art in correlation estimation.

Outperforms existing methods across all evaluated metrics on seven datasets.

Efficient end-to-end training with a transformer-based parameterization.

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.