[Paper Review] TABDDPM: Modeling Tabular Data with Diffusion Models

해당 논문에 대한 링크는 아래와 같다.

TabDDPM: Modelling Tabular Data with Diffusion Models

Denoising diffusion probabilistic models are currently becoming the leading paradigm of generative modeling for many important data modalities. Being the most prevalent in the computer vision community, diffusion models have also recently gained some attention in other domains, including speech, NLP, and graph-like data. In this work, we investigate if the framework of diffusion models can be advantageous for general tabular problems, where datapoints are typically represented by vectors of heterogeneous features. The inherent heterogeneity of tabular data makes it quite challenging for accurate modeling, since the individual features can be of completely different nature, i.e., some of them can be continuous and some of them can be discrete. To address such data types, we introduce TabDDPM -- a diffusion model that can be universally applied to any tabular dataset and handles any type of feature. We extensively evaluate TabDDPM on a wide set of benchmarks and demonstrate its superiority over existing GAN/VAE alternatives, which is consistent with the advantage of diffusion models in other fields. Additionally, we show that TabDDPM is eligible for privacy-oriented setups, where the original datapoints cannot be publicly shared.

arXiv.orgAkim Kotelnikov

TabDDPM의 공식 Code 구현은 아래의 링크에서 확인할 수 있다.

GitHub - yandex-research/tab-ddpm: [ICML 2023] The official implementation of the paper “TabDDPM: Modelling Tabular Data with Diffusion Models”

[ICML 2023] The official implementation of the paper “TabDDPM: Modelling Tabular Data with Diffusion Models” - yandex-research/tab-ddpm

GitHubyandex-research

0. Abstract

Denoisiong Diffusion Probabilistic Model (DDPM)은 많은 data modalities에 대하여 generative modeling의 paradigm을 주도하고 있다.

논문의 저자들은 General Tabular problem에 대하여 Diffusion Model의 framework가 유의미하게 작용할 수 있는지 알아보았다.

이때, general tabular problem들은 datapoint들이 heterogeneous feature들의 vector로 주로 표현된다.

각각의 individual feature들은 완전히 다른 nature를 갖고 있을 수 있기에(continuous, discrete) 이러한 tabular data의 heterogeneity는 정확한 modeling을 힘들게 한다.

이러한 data type들을 다루기 위해, 저자들은 TabDDPM을 소개한다.

이 Model은 보편적으로 어떠한 tabular dataset에도 적용가능하고, 어떠한 type의 feature도 다룰 수 있다.

저자들은 TabDDPM을 wide set of benchmark에 대하여 evaluate하고, 기존의 VAE/GAN 기반 tabular generative model에 대한 superiority를 주장한다.

추가적으로, TabDDPM은 original datapoint들이 publicly하게 공유될 수 없는 privacy-oriented setup에 대해서도 적용가능하다고 한다.

1. Introduction

DDPM (Denoising Diffusion Probabilistic Models)은 sampling하는 각 individual sample의 quality가 높으면서(realistic) diversity도 다양해서 다른 alternative approach들을 outperform했다.

그렇기에 DDPM은 generative modeling community에서 좋은 research interest의 대상이 되었다.

DDPM은 주로 Natural Image의 domain(Computer Vision)에서 impressive success들이 일어났는데, colorization, inpainting, segmentation, super-resolution, semantic editing 등과 같은 application분야에서 성공적으로 이용되었다.

위의 Image 뿐만 아니라 다른 Domain에서도 DDPM Framework들이 사용되었는데, 이에 NLP, Waveform signal processing, molecular graphs, time series 등으로 wide range of problem에 대해 적용되었다.

본 연구의 핵심은 DDPM의 universality가 general tabular problem까지 적용될 수 있는지에 있다.

여기서 general tabular problem은 set of heterogeneous feature를 가지는 dataset에 대해 industrial application을 할 때 보편적으로 발생하는 문제를 말한다.

즉, DDPM을 Tabular Data에 적용하려고 할 때 중요한 point는 크게 2가지가 있다:

1) Modern Privacy Regulations

High-quality generative model에 대한 수요가 급증하는데, 이는 GDPR과 같은 modern privacy regulation에 따라 real user data publishing을 방지하면서도 generative model이 생성한 synthetic data는 공유가 가능하기 때문이다.

2) Training High-quality Model of Tabular Data

CV나 NLP Domain에 비해 Tabular Data는 individual feature들의 heterogeneity와 상대적으로 작은 크기의 typical tabular datasets로 인해 High-quality model을 training시키기 힘들다.

이러한 2개의 문제에도 불구하고, 저자들은 tabular data의 typical distribution을 성공적으로 approximate할 수 있는 Diffusion Model을 제안한다.

TabDDPM의 Main Contribution은 총 3가지로 다음과 같다:

1. Introduce TabDDPM

categorical(discrete), numerical(continuous) feature들을 모두 포함하는 mixed data에 대해 작업가능하고, 어떠한 tabular task 에도 적용가능한 Simplest design of TabDDPM for tabular problem을 제안했다.

2. Outperformance

GAN-based, VAE-based model들을 포함하여 tabular data를 위해 design된 alternative approach들을 outperform한다. 또한 serveral datasets에 대한 advantage들을 설명한다.

3. Privacy concern

TabDDPM으로 생성한 tabular data들은 privacy concern scenario에 의해 real user data를 공유할 수 없을 때 "sweet spot"으로 작용할 수 있다.

3. Background

3.1) Diffusion Models

Diffusion Model은 크게 아래의 2 논문에 기반하여 발전했다.

Deep Unsupervised Learning using Nonequilibrium Thermodynamics

A central problem in machine learning involves modeling complex data-sets using highly flexible families of probability distributions in which learning, sampling, inference, and evaluation are still analytically or computationally tractable. Here, we develop an approach that simultaneously achieves both flexibility and tractability. The essential idea, inspired by non-equilibrium statistical physics, is to systematically and slowly destroy structure in a data distribution through an iterative forward diffusion process. We then learn a reverse diffusion process that restores structure in data, yielding a highly flexible and tractable generative model of the data. This approach allows us to rapidly learn, sample from, and evaluate probabilities in deep generative models with thousands of layers or time steps, as well as to compute conditional and posterior probabilities under the learned model. We additionally release an open source reference implementation of the algorithm.

arXiv.orgJascha Sohl-Dickstein

Denoising Diffusion Probabilistic Models

We present high quality image synthesis results using diffusion probabilistic models, a class of latent variable models inspired by considerations from nonequilibrium thermodynamics. Our best results are obtained by training on a weighted variational bound designed according to a novel connection between diffusion probabilistic models and denoising score matching with Langevin dynamics, and our models naturally admit a progressive lossy decompression scheme that can be interpreted as a generalization of autoregressive decoding. On the unconditional CIFAR10 dataset, we obtain an Inception score of 9.46 and a state-of-the-art FID score of 3.17. On 256x256 LSUN, we obtain sample quality similar to ProgressiveGAN. Our implementation is available at https://github.com/hojonathanho/diffusion

arXiv.orgJonathan Ho

둘 다 likelihood-based generative model이며 Forward Diffusion Process와 Reverse Denoising Processs 모두 Markov Process를 따른다.

Forward Process: Approximate Posterior $q(\mathbf{x}_{1:T}|\mathbf{x}_0)$

$$q(\mathbf{x}_{1:T}|\mathbf{x}_0) = \prod_{t=1}^T q(\mathbf{x}_t|\mathbf{x}_{t-1})$$

Initial sample $\mathbf{x}_0$에 각 timestep마다 점진적으로 scheduled된 noise를 주입한다.

즉, Initial data distribution $q(\mathbf{x}_0)$에서 시작하여 predefined distribution $q(\mathbf{x}_t|\mathbf{x}_{t-1})$ with variances ${\beta_1, ..., \beta_T}$으로부터 noise $\epsilon$를 sampling하여 gradually 더한다.

Reverse Process: Joint Distribution $p_{\theta}(\mathbf{x}_{0:T})$

$$p_{\theta}(\mathbf{x}_{0:T}) = \prod_{t=1}^T p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t)$$

time step $T$에서 시작하여 점진적으로 latent variable $\mathbf{x}_T \sim p(\mathbf{x}_T) = \mathcal{N}(\mathbf{x}_T; 0, \mathbf{I})$를 denoise한다.

결과적으로 $p_\theta(\mathbf{x}_0)$로부터 new data sample을 generate한다.

여기서 $q(\mathbf{x}_{t-1}|\mathbf{x}_{t})$는 intractable하기에, $p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t}) \approx q(\mathbf{x}_{t-1}|\mathbf{x}_{t})$를 만족하는 tractable한 $p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t})$를 도입한다.

결국 Reverse Denoising Process를 통해 $p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t}) \approx q(\mathbf{x}_{t-1}|\mathbf{x}_{t})$가 되도록 학습시킨다면,

subtract하는 gaussian noise를 예측할 수 있어 기존의 input data의 distribution과 유사한 distribution을 갖는 data를 생성할 수 있게 된다.

즉, $p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t})$의 분포가 $ q(\mathbf{x}_{t-1}|\mathbf{x}_{t})$에 approximate되도록 DDPM Model parameter $\theta$를 Neural Network를 통해 update시킨다.

이때, parameter $\theta$가 update되는 과정은 Variational Lower(Upper) Bound를 Optimize하면서 학습된다.

DDPM에서는 아래 $(3)$과 같이 유도된 Loss Function $L$을 바로 사용하지 않고, 한 번 더 형태를 변형하여 더 간결하게 $(5)$로 만든다.

$$ \mathbb{E}_{q(\mathbf{x}_0)}[-\log p_\theta(\mathbf{x}_0)] \leq \mathbb{E}_{q(\mathbf{x}_{0:T})} \left[ -\log \left( \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} | \mathbf{x}_0)} \right) \right] =\mathbb{E}_{q(\mathbf{x}_{0:T})} \left[ -\log p(\mathbf{x}_T) - \sum_{t=1}^T \log \frac{p_\theta(\mathbf{x}_{t-1} | \mathbf{x}_t)}{q(\mathbf{x}_t | \mathbf{x}_{t-1})} \right] := L \qquad(3) $$

$$L = \mathbb{E}_{q(\mathbf{x}_{0:T})}\left[D_{KL}(q(\mathbf{x}_T | \mathbf{x}_0)||p(\mathbf{x}_T)) + \sum_{t=2}^T D_{KL}(q(\mathbf{x}_{t-1} | \mathbf{x}_t, \mathbf{x}_0)||p_\theta(\mathbf{x}_{t-1} |\mathbf{x}_t))-\log p_\theta(\mathbf{x}_0 | \mathbf{x}_1)\right] \qquad (5)$$

3.2) Gaussian Diffusion Models

위와 마찬가지로 DDPM에 기반하여 설명한다.

Denoising Diffusion Probabilistic Models

We present high quality image synthesis results using diffusion probabilistic models, a class of latent variable models inspired by considerations from nonequilibrium thermodynamics. Our best results are obtained by training on a weighted variational bound designed according to a novel connection between diffusion probabilistic models and denoising score matching with Langevin dynamics, and our models naturally admit a progressive lossy decompression scheme that can be interpreted as a generalization of autoregressive decoding. On the unconditional CIFAR10 dataset, we obtain an Inception score of 9.46 and a state-of-the-art FID score of 3.17. On 256x256 LSUN, we obtain sample quality similar to ProgressiveGAN. Our implementation is available at https://github.com/hojonathanho/diffusion

arXiv.orgJonathan Ho

위의 3.1) Diffusion Models에서 forward / reverse process가 모두 Gaussian Distribution으로 characterize될 경우 Loss Function은 다음과 같이 $L_t^{simple}$로 표현가능하다.

Forward Diffusion Process $q(\mathbf{x}_t|\mathbf{x}_{t-1})$와 Reverse Denoising Process $q(\mathbf{x}_{t-1}|\mathbf{x}_{t})$들이 Gaussian Distribution으로 characterized 되었다.

그리고 이러한 Forward/Reverse process가 정의된 Gaussian Diffusion Model은 Continuous space $\mathbf{x}_t \in \mathbb{R}^n$에서 작동한다.

• Forward Diffusion Process: $q(\mathbf{x}_t|\mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{1-\beta_t}\mathbf{x}_{t-1}, \beta_t I)$
• $p(\mathbf{x}_T) = \mathcal{N}(\mathbf{x}_T; 0, \mathbf{I})$
• Reverse Denoising Process: $p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \mu_\theta(\mathbf{x}_t,t), \Sigma_\theta(\mathbf{x}_t,t))$

1) 이때, DDPM은 $\Sigma_{\theta}(\mathbf{x}_{t}, t))$를 $\sigma_t^{2}I$로 설정해서 time-dependent constant로 처리하도록 한다.

$$\Sigma_{\theta}(\mathbf{x}_{t}, t)) = \sigma_t^{2}I$$

$$\sigma_t^{2} = \beta_t \: \text{or} \: \sigma_t^{2} = \tilde{\beta_t}$$


$$\tilde{\beta}_t := \frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t}\beta_t$$

위의 2가지 case 중 어떠한 것을 골라도 실험적으로 비슷한 결과가 나온다고 한다.

2) $\mu_{\theta}(\mathbf{x}_t, t)$는 특정 parameterization을 이용하여 $p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t}) = \mathcal{N}(\mathbf{x}_{t-1}; \mu_{\theta}(\mathbf{x}_t, t), \sigma_t^{2})$에 대해 다음과 같이 작성할 수 있다.

$$L_{t-1} = \mathbb{E}_q\left[\frac{1}{2\sigma_t^2} ||\tilde{\mu_t}(\mathbf{x}_t, \mathbf{x}_0) - \mu_\theta(\mathbf{x}_t, t)||^2\right] + C$$

• $C$는 parameter $\theta$에 의존하지 않는 constant이다.

따라서 $L_{t-1}$은 learnable parameter $\mu_\theta(\mathbf{x}_t, t)$가 ground truth $\tilde{\mu_t}(\mathbf{x}_t, \mathbf{x}_0)$에 가까워지도록 학습한다.

$$L_{t-1} = \mathbb{E}_q\left[\frac{1}{2\sigma_t^2} ||\tilde{\mu_t}(\mathbf{x}_t, \mathbf{x}_0) - \mu_\theta(\mathbf{x}_t, t)||^2\right] + C$$

3) 위의 Loss는 추가적으로 Paramterization이 가능하다.

$\mathbf{x}_t$를 reparameterization한 식에서 $\epsilon$에 대해 정리해보자.

$$\mathbf{x}_t(\mathbf{x}_0, \epsilon) = \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\epsilon, \quad \epsilon \sim \mathcal{N}(0,I)$$

$$\mathbf{x}_0 = \frac{1}{\sqrt{\bar{\alpha}_t}}(\mathbf{x}_t(\mathbf{x}_0,\epsilon) - \sqrt{1 - \bar{\alpha}_t}\epsilon)$$

$\mathbf{x}_0$를 $\tilde{\mu}_t(\mathbf{x}_t,\mathbf{x}_0)$의 식에 대입하여 정리하자.

$$q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; \tilde{\mu}_t(\mathbf{x}_t,\mathbf{x}_0), \tilde{\beta}_t I)$$


$$\tilde{\mu}_t(\mathbf{x}_t,\mathbf{x}_0) := \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}\mathbf{x}_0 + \frac{\sqrt{1-\bar{\beta}_t}(1-\bar{\alpha} _{t-1})}{1-\bar{\alpha}_t}\mathbf{x}_t \quad \text{and} \quad \tilde{\beta}_t := \frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t}\beta_t$$

$$\tilde{\mu}_t(\mathbf{x}_t, \mathbf{x}_0) = \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}\mathbf{x}_0 + \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}\mathbf{x}_t $$

$$= \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}\frac{1}{\sqrt{\bar{\alpha}_t}}(\mathbf{x}_t(\mathbf{x}_0,\epsilon) - \sqrt{1 - \bar{\alpha}_t}\epsilon) + \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}\mathbf{x}_t $$

$$= \left(\frac{\sqrt{\bar{\alpha}_{t-1}}}{\sqrt{\bar{\alpha}_t}}\frac{\beta_t}{1-\bar{\alpha}_t} + \frac{1}{\sqrt{\alpha_t}}\frac{\alpha_t(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}\right)\mathbf{x}_t(\mathbf{x}_0,\epsilon) - \frac{\sqrt{\bar{\alpha}_{t-1}}}{\sqrt{\bar{\alpha}_t}}\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon $$

$$= \frac{1}{\sqrt{\alpha_t}}\left(\left(\frac{\beta_t + \alpha_t - \alpha_t\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t}\right)\mathbf{x}_t(\mathbf{x}_0,\epsilon) - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon\right) $$

$$= \frac{1}{\sqrt{\alpha}_t}\left(\mathbf{x}_t(\mathbf{x}_0,\epsilon) - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon\right)$$

$$L_{t-1} - C = \mathbb{E}_{\mathbf{x}_0,\epsilon}\left[\frac{1}{2\sigma_t^2}\left\|\frac{1}{\sqrt{\alpha}_t}\left(\mathbf{x}_t(\mathbf{x}_0,\epsilon) - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon\right) - \mu_\theta(\mathbf{x}_t(\mathbf{x}_0,\epsilon),t)\right\|^2\right]$$

결국 learnable $\mu_\theta(\mathbf{x}_t(\mathbf{x}_0,\epsilon),t)$가 주어진 ground truth $\tilde{\mu}_t(\mathbf{x}_t, \mathbf{x}_0) $인 $\mathbf{x}$와 $\epsilon$으로만 이루어진 $\frac{1}{\sqrt{\alpha}_t}\left(\mathbf{x}_t(\mathbf{x}_0,\epsilon) - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon\right)$을 예측하도록 학습해야 한다.

그런데, 주어진 time step $t$에 대하여 time-dependent한 constant가 아닌 항은 $\epsilon$이 유일하므로 결국 $\mu_\theta(\mathbf{x}_t(\mathbf{x}_0,\epsilon),t)$는 $\epsilon_{\theta}(\mathbf{x}_t, t)$을 예측하는 것으로 귀결된다.

따라서 $\epsilon_{\theta}(\mathbf{x}_t, t)$에 대해 식을 전개해야 하는데, $\mu_\theta(\mathbf{x}_t(\mathbf{x}_0,\epsilon),t)$가 $\tilde{\mu}_t(\mathbf{x}_t, \mathbf{x}_0) $로 가까워질 때 식의 구조 자체는 동일하되, $\epsilon$만 learnable parameter $\theta$를 가지고 있는 $\epsilon_\theta(\mathbf{x}_t, t)$로 바뀐다.

$$\mu_\theta(\mathbf{x}_t(\mathbf{x}_0,\epsilon),t) = \tilde{\mu}_t(\mathbf{x}_t, \frac{1}{\sqrt{\bar{\alpha}_t}}(\mathbf{x}_t(\mathbf{x}_0,\epsilon) - \sqrt{1 - \bar{\alpha}_t}\epsilon_{\theta}(\mathbf{x}_t, t))) = \frac{1}{\sqrt{\alpha}_t}\left(\mathbf{x}_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_{\theta}(\mathbf{x}_t, t)\right)$$

learnable parameter $\mu_\theta(\mathbf{x}_t(\mathbf{x}_0,\epsilon),t)$안의 $\epsilon$가 learnable 하기에 $\epsilon_{\theta}(\mathbf{x}_t, t)$로 취급한다.

$$= \mathbb{E}_{\mathbf{x}_0,\epsilon} \left[ \frac{\beta^2}{2\sigma^2_t\alpha_t(1 - \bar{\alpha_t})} ||\epsilon - \epsilon_\theta(\sqrt{\bar{\alpha_t}}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha_t}}\epsilon, t)||^2 \right]$$

• Learnable Parameter: $$\mu_\theta(\mathbf{x}_t(\mathbf{x}_0,\epsilon),t) = \frac{1}{\sqrt{\alpha}_t}\left(\mathbf{x}_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_{\theta}(\mathbf{x}_t, t)\right)$$

• Ground Truth: $$\tilde{\mu}_t(\mathbf{x}_t, \mathbf{x}_0) = \frac{1}{\sqrt{\alpha}_t}\left(\mathbf{x}_t(\mathbf{x}_0,\epsilon) - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon\right)$$

위의 두 parameter $\mu_\theta(\mathbf{x}_t(\mathbf{x}_0,\epsilon),t) , \tilde{\mu}_t(\mathbf{x}_t, \mathbf{x}_0)$를 아래 Loss $L_{t-1}$에 대입해서 정리하면 다음과 같이 간단하게 변환될 수 있다.

$$L_{t-1} = \mathbb{E}_q\left[\frac{1}{2\sigma_t^2} ||\tilde{\mu_t}(\mathbf{x}_t, \mathbf{x}_0) - \mu_\theta(\mathbf{x}_t, t)||^2\right] + C$$

$$=\mathbb{E}_{\mathbf{x}_0,\epsilon} \left[ \frac{\beta^2}{2\sigma^2_t\alpha_t(1 - \bar{\alpha_t})} ||\epsilon - \epsilon_\theta(\sqrt{\bar{\alpha_t}}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha_t}}\epsilon, t)||^2 \right]$$

따라서, 위의 objective function은 아래와 같이 fixed noise $\epsilon$과 learnable noise $\epsilon_\theta(\mathbf{x}_t, t)$ 간의 MSE(Mean-Squared Error)로 simplify될 수 있다.

$$L_t^{simple} = \mathbb{E}_{\mathbf{x}_0,\epsilon,t} ||\epsilon - \epsilon_\theta(\mathbf{x}_t,t)||_2^2$$

3.3) Multinomial Diffusion Models

Multinomial Diffusion Model에 대한 상세한 설명은 아래의 링크에서 확인할 수 있다.

Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions

Generative flows and diffusion models have been predominantly trained on ordinal data, for example natural images. This paper introduces two extensions of flows and diffusion for categorical data such as language or image segmentation: Argmax Flows and Multinomial Diffusion. Argmax Flows are defined by a composition of a continuous distribution (such as a normalizing flow), and an argmax function. To optimize this model, we learn a probabilistic inverse for the argmax that lifts the categorical data to a continuous space. Multinomial Diffusion gradually adds categorical noise in a diffusion process, for which the generative denoising process is learned. We demonstrate that our method outperforms existing dequantization approaches on text modelling and modelling on image segmentation maps in log-likelihood.

arXiv.orgEmiel Hoogeboom

Multinomial Diffusion Model은 categorical data(discrete)인 $x_t \in \left \{0, 1\right \}^K$를 생성하기 위해 design된 모델이다.

$x_t \in \left \{0, 1\right \}^K$은 K개의 value(원소)를 가지는 one-hot encoded categorical variable 이다.

Multinomial Forward Diffusion Process는 $K$ class에 대한 Uniform Noise로부터 corrupt 시키는 categorical distribution $q(\mathbf{x}_t|\mathbf{x}_{t-1})$이다.

• $q(\mathbf{x}_t|\mathbf{x}_{t-1}) := Cat(\mathbf{x}_t; (1-\beta_t)\mathbf{x}_{t-1} + \beta_t/K)$
• $q(\mathbf{x}_T) := Cat(\mathbf{x}_T; 1/K)$
• $q(\mathbf{x}_t|\mathbf{x}_0) = Cat(\mathbf{x}_t; \bar{\alpha_t} \mathbf{x}_0 + (1-\bar{\alpha_t})/K)$

$Cat$ denotes a categorical distribution with probability parameters after ;.

위의 식으로부터 posterior $q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0)$ 를 유도할 수 있다.

$$q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0) = Cat(\mathbf{x}_{t-1}; \pi / \sum_{k=1}^K \pi_k)$$


where $\pi = [\alpha_t \mathbf{x}_t + (1-\alpha_t)/K] \circ [\bar{\alpha_{t-1}}\mathbf{x}_0 + (1-\bar{\alpha_{t-1}})/K]$.

Multinomial Reverse Denoising Process $p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)$는 $q(\mathbf{x}_{t-1}|\mathbf{x}_t,\hat{\mathbf{x}}_0(\mathbf{x}_t,t))$으로 parameterize된다.

이때, $\hat{\mathbf{x}}_0$는 Neural Network를 통해 학습된다.

이후, Multinomial Diffusion Model은 Variational Lower Bound를 Maximize하면서 Training을 진행한다.

4. TabDDPM

이번 Section 4에서는 TabDDPM의 Design과 함께 Model's Effectiveness에 영향을 주는 Main hyperparameter를 함께 살펴볼 것이다.

4.1) TabDDPM Design

TabDDPM은 Categorical Feature(Discrete) / Numerical Feature(Continuous) 각각에 대하여 사용하는 Diffusion Model이 다르다.

• binary feature와 categorical feature: Multinomial Diffusion 사용
• numerical feature: Gaussian Diffusion 사용

Tabular Data sample $x = \left[ x_{\text{num}}, x_{\text{cat}_1}, \cdots, x_{\text{cat}_C} \right]$은 다음의 특징들을 가진다.

- Numerical Data $x_{\text{num}}$: $N_{\text{num}}$개의 numerical feature를 가지며, $x_{\text{num}} \in \mathbb{R}^{N_{\text{num}}}$



- Categorical Data $x_{\text{cat}_i}$: $C$개의 categorical feature를 가지며, 각 $i$번째 feature는 $K_i$개의 category(class)를 가진다.

왜 굳이 $N_{\text{num}}$개의 Numerical Feature를 가지는 column vector를 $x_{\text{num}}$ 하나만 넣는가?


각 categorical feature인 $x_{\text{cat}_i}$는 총 C개가 있지 않은가?

$x = \left[ x_{\text{num}}, x_{\text{cat}_1}, \cdots, x_{\text{cat}_C} \right]$를 전체 Dataset이 아니라 하나의 Column으로 봐야 한다.

즉, 하나의 Column $x$에 $N_{\text{num}}$개의 Numerical Feature들이 있고, $C$개의 Categorical Feature들이 있는데, 각 $i$th Categorical Feature는 각각 $K_i$개의 class 중 하나의 category를 가진다.

시작에서부터 언급했듯이, TabDDPM은 하나의 Column 안에 numerical, categorical feature들이 모두 섞여있는 mixed data에 대하여 작동할 수 있도록 하는 것이 목적이다.

따라서 해당 TabDDPM의 Design은 Dataset의 각 $N_{\text{sum}} + C$개의 element를 가지는 Column을 하나의 단위로 인식하여 Processing한 것이다.

• Numercial Feature: Gaussian Quantile Transformation (scikit-learn library)를 통해 normalize 시켜 $x_{\text{num}}^{\text{norm}}$으로 변환한다.

• Categorical Feature: One-Hot encoding을 통해 $i$번째 feature $x_{\text{cat}_i}$를 $x_{\text{cat}_i}^{\text{ohe}} \in \left\{ 0, 1\right\}^{K_i}$로 변환한다.

따라서, 최종 Input $x_0$는 $N_{\text{num}} + \Sigma K_i$의 dimensionality를 가진다.

Forward Diffusion Process: 각 Categorical feature는 개별적인 forward process에 의해 handling된다.

즉, 각 feature들에 대한 noise component들은 독립적으로 sampling된다.

Reverse Desnoising Process: TabDDPM에서 사용되는 reverse process는 Multi-Layer Network를 사용한다.

Reverse diffusion step에서의 MLP는 input $x_0$와 동일한 dimensionality를 가지는 output을 출력한다.

output의 first $N_{\text{num}}$ coordinate는 Gaussian Diffusion에 의한 $\epsilon$의 prediction이고,


output의 나머지는 Multinomial Diffusion에 의한 $x_{\text{cat}_i}^{\text{ohe}}$의 prediction이다.

위의 Figure는 TabDDPM을 Classification Problem에 적용하기 위한 Design 이다.

TabDDPM Loss는 2개의 Term으로 구성되어 있는데, Gaussian Diffusion Term(numerical)인 $L_t^{\text{simple}}$과 Multinomial Diffusion Term (categorical)인 $L_t^i$들로 구성되어 있다.

• $L_t^{\text{simple}}$: Gaussian Diffusion Term (numerical, continuous)
  • Mean-Squared Error (MSE) of gaussian noise $\epsilon$

$$L_t^{simple} = \mathbb{E}_{\mathbf{x}_0,\epsilon,t} ||\epsilon - \epsilon_\theta(\mathbf{x}_t,t)||_2^2$$

• $L_t^i$: Multinomial Diffusion Term (Categorical, Discrete)
  • KL Divergence for each $i$th feature

$$L_t = D_{KL}(q(\mathbf{x}_{t-1} | \mathbf{x}_t, \mathbf{x}_0)||p_\theta(\mathbf{x}_{t-1} |\mathbf{x}_t))$$

그리고 Total Loss의 Multinomial diffusion term은 각 Categorical Feature 마다 더하고 $\sum_{i\leq C} L_t^i}$, Categorical Feature 의 개수 $C$로 나눈다.

$$L_t^{TabDDPM} = L_t^{simple} + \frac{\sum_{i\leq C} L_t^i}{C}$$

왜 Main Structure를 U-Net으로 하지 않고, MLP를 사용했나?

저자들은 아래의 논문에서 제시한 구조를 Adapt 했다고 한다.

왜 굳이 U-Net이 아니라 MLP를 사용했는지에 관해서는 아래의 논문을 읽어보자.

(아마 Tabular Data에 fit하도록 Modeling을 하기 위해서 사용한 것 같기는 하다)

Revisiting Deep Learning Models for Tabular Data

The existing literature on deep learning for tabular data proposes a wide range of novel architectures and reports competitive results on various datasets. However, the proposed models are usually not properly compared to each other and existing works often use different benchmarks and experiment protocols. As a result, it is unclear for both researchers and practitioners what models perform best. Additionally, the field still lacks effective baselines, that is, the easy-to-use models that provide competitive performance across different problems. In this work, we perform an overview of the main families of DL architectures for tabular data and raise the bar of baselines in tabular DL by identifying two simple and powerful deep architectures. The first one is a ResNet-like architecture which turns out to be a strong baseline that is often missing in prior works. The second model is our simple adaptation of the Transformer architecture for tabular data, which outperforms other solutions on most tasks. Both models are compared to many existing architectures on a diverse set of tasks under the same training and tuning protocols. We also compare the best DL models with Gradient Boosted Decision Trees and conclude that there is still no universally superior solution.

arXiv.orgYury Gorishniy

Classification Dataset을 위해서는, $p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t}, \mathbf{y})$를 학습하는 class-conditional model을 사용한다.

Regression Dataset을 위해서는, target value를 추가적인 numerical feature로 취급하여 joint distribution을 학습한다.

Reverse Denoising Process를 modeling하기 위해서는, 아래의 논문에서 사용된 MLP architecture를 채택한다.

Revisiting Deep Learning Models for Tabular Data

The existing literature on deep learning for tabular data proposes a wide range of novel architectures and reports competitive results on various datasets. However, the proposed models are usually not properly compared to each other and existing works often use different benchmarks and experiment protocols. As a result, it is unclear for both researchers and practitioners what models perform best. Additionally, the field still lacks effective baselines, that is, the easy-to-use models that provide competitive performance across different problems. In this work, we perform an overview of the main families of DL architectures for tabular data and raise the bar of baselines in tabular DL by identifying two simple and powerful deep architectures. The first one is a ResNet-like architecture which turns out to be a strong baseline that is often missing in prior works. The second model is our simple adaptation of the Transformer architecture for tabular data, which outperforms other solutions on most tasks. Both models are compared to many existing architectures on a diverse set of tasks under the same training and tuning protocols. We also compare the best DL models with Gradient Boosted Decision Trees and conclude that there is still no universally superior solution.

arXiv.orgYury Gorishniy

• tabular input $x_{\text{in}}$: $\left [x_{\text{num}}^{\text{norm}}, x_{\text{cat}_1}^{\text{ohe}}, ... , x_{\text{cat}_C}^{\text{ohe}}\right]$
• time step $t$
• class label $y$

$t_{\text{emb}} = \text{Linear}(\text{SiLU}(\text{Linear}(\text{SinTimeEmb}(t))))$


$y_{\text{emb}} = \text{Embedding}(y)$


$x = \text{Linear}(x_{\text{in}}) + t_{\text{emb}} + y_{\text{emb}}$


여기서 사용하는 모든 $\text{Linear}$ Layer들은 128이라는 고정된 dimension을 사용한다.

따라서 아래와 같이 MLP Layer를 쌓을 수 있다.

$$\text{MLPBlock}(x) = \text{Dropout}(\text{ReLU}(\text{Linear}(x)))$$

$$\text{MLP}(x) = \text{Linear}(\text{MLPBlock}(...(\text{MLPBlock}(x))))$$

여기서 사용하는 $\text{SinTimeEmb}$는 Sinusodial Time Embedding으로 128 dimension을 사용하며, Transformer에서 사용하는 Positional Encoding으로 생각하면 된다.

4.2) TabDDPM Hyperparameters

TabDDPM에서 사용하는 Parameter들은 Model Design과 마찬가지로 model effectiveness에 강한 영향을 주기에 필수적으로 고려되어야 한다.

위의 Table 1은 Hyperparameter들을 명시함과 동시에 이들의 Search Space까지 보여주고 있다.

자세한 Hyperparameter tuning의 process는 Experimental section에서 자세하게 설명하고 있다.