[Paper Review] Mixed-Type Tabular Data Synthesis with Score-based Diffusion in Latent Space

Hengrui Zhang, Jiani Zhang+6 more|arXiv (Cornell University)|Oct 14, 2023

Cancer-related molecular mechanisms research23 citations

TL;DR

TabSyn proposes a diffusion model in a learned latent space (via a VAE with Transformers and adaptive loss) to synthesize mixed-type tabular data, achieving faster sampling and higher data quality than prior methods.

ABSTRACT

Recent advances in tabular data generation have greatly enhanced synthetic data quality. However, extending diffusion models to tabular data is challenging due to the intricately varied distributions and a blend of data types of tabular data. This paper introduces Tabsyn, a methodology that synthesizes tabular data by leveraging a diffusion model within a variational autoencoder (VAE) crafted latent space. The key advantages of the proposed Tabsyn include (1) Generality: the ability to handle a broad spectrum of data types by converting them into a single unified space and explicitly capture inter-column relations; (2) Quality: optimizing the distribution of latent embeddings to enhance the subsequent training of diffusion models, which helps generate high-quality synthetic data, (3) Speed: much fewer number of reverse steps and faster synthesis speed than existing diffusion-based methods. Extensive experiments on six datasets with five metrics demonstrate that Tabsyn outperforms existing methods. Specifically, it reduces the error rates by 86% and 67% for column-wise distribution and pair-wise column correlation estimations compared with the most competitive baselines.

Motivation & Objective

• Motivate improved generative modeling for mixed-type tabular data with complex inter-column relationships.
• Develop a unified latent space where numerical and categorical features can be jointly modeled.
• Improve sample quality and generation speed relative to existing diffusion-based tabular data methods.
• Demonstrate robustness across multiple datasets and evaluation metrics including low- and high-order statistics and downstream tasks.

Proposed method

• Transform raw tabular data into a continuous embedding space via column-wise tokenizers and Transformer encoders.
• Use a variational autoencoder with adaptive loss weighting to learn informative, smoothed latent embeddings.
• Train a score-based diffusion model in the latent space using denoising score matching with a linear noise schedule (sigma(t)=t).
• Apply a detokenizer to reconstruct original column values from latent representations for data generation.
• Provide theoretical justification and empirical evidence that linear noise schedules reduce reverse-process errors and accelerate sampling (fewer than 20 NFEs).
• Evaluate TabSyn against seven baselines on six mixed-type datasets using column-wise density, pairwise correlations, MLE, and missing-value imputation tasks.

Figure 1: TabSyn wins: Our TabSyn method consistently outperforms SOTA tabular data generation methods across five data quality metrics.

Experimental results

Research questions

• RQ1 Can a diffusion model operate effectively in a latent space learned from mixed-type tabular data?
• RQ2 Do column-aware tokenizers and a Transformer-based VAE better preserve inter-column relationships than existing methods?
• RQ3 Does a linear noise schedule enable fast and high-quality diffusion-based tabular data synthesis?
• RQ4 How does TabSyn perform across low-order (density, correlation) and high-order (precision/recall) metrics and downstream tasks?
• RQ5 Is adaptive beta-VAE training beneficial for latent space regularization and reconstruction quality?

Key findings

• TabSyn consistently outperforms baseline methods on column-wise density estimation and pairwise column correlations, reducing error by average 86.0% and 67.6% respectively.
• TabSyn achieves superior machine learning efficiency (classification/regression) and missing-value imputation performance across six datasets.
• A linear diffusion noise schedule (sigma(t)=t) enables generating high-quality samples with fewer than 20 reverse steps, accelerating sampling.
• Latent-space diffusion (TabSyn-DDPM) generally outperforms diffusion in data space, highlighting benefits of learning latent embeddings for tabular data.
• A scheduled beta-VAE improves reconstruction without over-regularizing the latent space, enhancing downstream synthetic data quality.

Figure 2: An overview of the proposed TabSyn . Each row data $x$ is mapped to latent space $z$ via a column-wise tokenizer and an encoder. A diffusion process $z_{0}\rightarrow z_{T}$ is applied in the latent space. Synthesis $z_{T}\rightarrow z_{0}$ starts from the base distribution $p(z_{T})$ and