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

Thank you for your constructive feedback and insightful questions! Below, we provide detailed responses to each of your questions. Additionally, please refer to our revised paper, with updates colored in purple.

[W1] Modeling feature dependencies

Thank you for raising the point regarding the conditional independence assumption made when estimating the posterior $q(x_s|x_t,x_0)$.

Factorization does not limit learning inter-feature dependencies. We first clarify that this assumption allows for easy optimization while still allowing the model to capture complex dependencies between different features. In the diffusion model, it’s a common practice to factorize the posterior estimation across all data dimensions. In practice, we not only factorize between numerical and categorical features but factorize across all individual features, in the same fashion as image diffusion models factorize across all pixels in the image grids. The detailed formulation is shown below:

This factorization does not imply that the numerical and categorical features are independent in the reverse generative process nor dependencies are ignored. This is because each conditional distribution $q(x_s^{\text{num}} \mid x_t, x_0)$ and $q(x_s^{\text{cat}} \mid x_t, x_0)$ is conditioned on the entire $x_t$ and $x_0$, which contains information from both numerical and categorical features. This allows the model to learn and capture complex dependencies between different features.

The efficacy of this approach in capturing cross-feature dependencies has been proven across many applications of diffusion models. Examples include modeling local spatial dependencies in image generation [1,2] and modeling the dependency between atom coordinates and types in molecular generation [3]. In the context of tabular data generation, this assumption is also effective in capturing dependencies between features of different types, as evidenced by its adoption in methods like TabDDPM, TabSyn, and ours.

How TabDiff’s design tricks further improve dependency modeling. Beyond this standard practice, TabDiff introduces feature-wise adaptive noise schedules that further enhance the model’s ability to capture dependencies between features. These learnable schedules give the model the flexibility to denoise different features at varying speeds. Specifically, a more positive $\rho$ and $k$ corresponds to quicker denoising. During training, the model learns different denoising rates for different features feature, which can be regarded as an order of generation. Features with larger $\rho$ and $k$ are recovered earlier in the denoising process, allowing subsequent features to be generated and conditioned on them. This mechanism enables the model to better capture inter–feature dependencies. In Table 2, TabDiff outperforms all other baselines on inter-column correlation metrics (i.e., Trend); in Table 5, TabDiff with learnable schedules achieves higher scores on Trend.

[1] Ho et al., 2020. “Denoising Diffusion Probabilistic Models.”

[2] Karras et al., 2022. “Elucidating the Design Space of Diffusion-Based Generative Models.”

[3] Hoogeboom et al., 2022. “Elucidating the Design Space of Diffusion-Based Generative Models.”

Thank you for your constructive feedback and insightful questions! Below, we provide detailed responses to each of your questions. Additionally, please refer to our revised paper, with updates colored in purple.

[W1] Model Efficiency

This is an excellent suggestion, and we totally agree that a thorough analysis of model efficiency could further enhance our methods’ evaluation. Below, we first compare the training and sampling speed of TabDiff against other competitive baseline methods (TabDDPM, TabSyn) using four different metrics. Additionally, inspired by your suggestion, we conduct a new ablation study exploring the tradeoff between discretization steps (i.e., efficiency) and sample quality in diffusion-based models. We use the representative Adult dataset, which contains a balanced number of numerical and categorical columns, throughout the experiment Our results show that, among the competitive methods, TabDiff is not only the most effective but also highly efficient. TabDiff also maintains the most capacity when the sampling discretization steps are limited. Below, we present detailed analyses.

Metric#1: Training Time. We first measure the training time of each method. The epoch lengths are set to the default configuration of each method, and the validation frequencies are set to the same value of once every 200 epochs. Our measurements are presented in the first column of the table below, with the entry of TabSyn being split into – VAE training time + Diffusion training time.

We see that all three methods take a comparable time to complete one training run, with Tabsyn being $10\%$ faster than TabDiff and TabDiff being $20\%$ faster than TabDDPM. The current gap between TabDiff and TabSyn is likely due to TabDiff’s slightly deeper network architecture compared to TabSyn (note: deeper but not larger, as we enforce the same parameter size for all models in our experiments). We believe that the architecture of TabDiff can be further optimized to improve efficiency, and we leave this as future work.

Nevertheless, it is also important to consider model robustness when assessing efficiency. As highlighted in our response to [W2] raised by Reviewer gKXC, the training process for TabSyn is notably unstable due to the difficulty in training VAEs, often requiring you to retrain many times in order to produce a model capable of generating samples comparable in quality to TabDiff. On the other hand, TabDDPM fails to scale to more complicated datasets as mentioned in the response to [W2] you raised. Thus, when taking into account training robustness, TabDiff is the most robust and efficient model among all competitive methods.

Metric#2: Training Convergence. Next, we assess training convergence by evaluating the quality of samples generated from intermediate checkpoints during the training process. Figure 6 in the revised manuscript plots our result. The curves show that TabDiff converges faster than the other methods, as TabDiff produces more high-quality samples when shown to the datasets for the same number of times (i.e., same epoch).

Metric#3: Number of Function Evaluation (NFE). For sampling, we first theoretically analyze the number of network calls involved in a single diffusion step (i.e. the denoising step from $x_t$ to $x_{t-1}$)) for each method. The numbers are presented in the third column of the table below. TabDiff and TabSyn involve an extra network call because TabDiff employs the second-order correction trick introduced in [1].

Metric#4: Sampling Time. We measure empirically the time to generate the same number of samples as the test set (32561 samples for Adult). The numbers are presented in the second column of the table below. According to it, TabDiff and TabSyn samples $\sim10\times$ faster than TabDDPM. This result is expected, as both TabDiff and TabSyn, by default, use 20 times fewer sampling steps than TabDDPM, while making twice as many network calls per step compared to TabDDPM. TabDiff’s slightly longer sampling time is also attributed to the evaluation of its deeper network. We believe this can be optimized in future work.

Dear Reviewer kzHM,

We would like to express our gratitude for your valuable comments and suggestions. Thanks a lot for your positive feedback on our method and experiments.

As the author-reviewer discussion period is coming to a close on Nov 27th, we would greatly appreciate it if you could clarify any remaining concerns you may have. This will ensure that we can adequately address them during this period.

Thank you very much for your attention to this matter.

Best regards,

Authors

[W2, Q2] Motivation of masked diffusion This is another good point. Our choice is primarily motivated by the following advantages that masked diffusion offers over other discrete diffusion frameworks.

1. Better numerical stability and efficiency in training. The masked diffusion framework has a very simple objective. Unlike more generalized discrete diffusion frameworks [1, 2], where the loss is evaluated across all columns, our framework computes the loss solely on columns that have been masked [Eq. (9)]. Moreover, our loss is weighted by the simple factor of $\frac{\alpha_t^{\prime}}{1-\alpha_t}$, which simplifies to $\frac{-k}{1-\alpha_t}$ under our noise schedule parameterization. This simple training objective enhances both training numerical stability and efficiency [3, 4].

These benefits were initially demonstrated in experiments on texts [3, 4]. Recognizing that categorical columns of tabular data can be considered sequences of unordered tokens similar to text, we were motivated to apply the simple masked diffusion framework to tabular data generation. As expected, we observed similar benefits of stable training.

To illustrate these, we draw a direct comparison with TabDDPM, which utilizes a general discrete diffusion framework that interpolates between data and a uniform prior. TabDiff shows more stable training when scaling up to datasets with more complex categorical columns. For instance, on the Diabetes datasets which include multiple categorical features with very large numbers of classes (Table 6), the training of TabDDPM failed to converge properly, as evidenced by the exploding training loss in Figure 10 of the revised paper. Consequently, TabDDPM failed to generate meaningful content, resulting in poor Shape and Trend scores as presented in Tables 1 and 2. In contrast, TabDiff, utilizing the simple masked diffusion framework, exhibits very stable training across all datasets and consistently produces high-quality samples. Furthermore, TabDiff is more efficient to train than TabDDPM, as shown by its shorter training time and faster convergence in Table 11 and Figure 6 of the revised paper.

2. Compatibility with Auxiliary Techniques. The second advantage is that the simple masked diffusion framework is easily compatible with auxiliary techniques [4]. In our paper, we successfully adopted the classifier-free guidance technique for conditional generation, in a fairly clean matter. In the Appendix B.2 revised version of the paper, we have shown that the CFG guided posterior probability $\tilde{p}\_{\theta}(x\_s^{\text{cat}} | x\_t, y)$ can be simply computed by directly interpolating the model's raw estimates of $x_0$, i.e.,$\mu_{\theta}^{\text{cat}}(x_t, y, t)$ and $\mu_{\theta}^{\text{cat}}(x_t, t)$.

[1] Austin et al., 2021. “Structured Denoising Diffusion Models in Discrete State-Spaces.”

[2] Campbel et al., 2022 “A Continuous Time Framework for Discrete Denoising Models.”

[3] Lou et al., 2024, “Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution.”

[4] Sahoo et al., 2024, “Simple and effective masked diffusion language models.”

Thank you for your constructive feedback and insightful questions! Below, we provide detailed responses to each of your questions. Additionally, please refer to our revised paper, with updates colored in purple.

[W1] How TabDiff tackles feature-wise heterogeneity

This is a very good point. We provide a more detailed discussion in the following:

Recap: what do we mean by feature-wise heterogeneity. By feature-wise heterogeneity, we refer to the fact that, in tabular data, each feature/column often exhibits marginal distribution in distinct shapes (which is illustrated by the blue density curves in Figure 11 of the revised paper). This poses an extra challenge for generative modeling: columns with more complex marginals are more difficult to model, so generative models should pay more attention to them.

How TabDiff handles this problem. In our approach, we address this challenge by implementing adaptive learnable noise schedules, as more thoroughly discussed in Section 2.3. These schedules allow the model to adaptively adjust the diffusion process for each feature throughout the training process. This, in turn, eventually lets the model optimally allocate its capacity to capture distinct distributions of different columns, reducing learning bias towards any specific column.

Evidence of the improvement. In our original submission, we demonstrated the effectiveness of this approach in our ablation study (Section 4.4, Table 5). Specifically, TabDiff with learnable schedules achieves a lower average error of column-wise density (i.e. Shape) across columns compared to a version with uniform schedules, indicating a better ability to handle heterogeneous feature distributions.

Addressing the Lack of Detailed Explanation. We totally agree with your point that we need to address what we mean by “optimal allocation of model capacity across heterogeneous features” in greater detail. Inspired by your feedback, we further study the impact of learnable schedules by analyzing how generation quality varies between different features. The Shape and Trend errors presented in Table 5 are averaged over all columns and columns pars. Now, in the revised manuscript, using the representative Adult dataset, we visualize the errors at each column and column pair in Figures 8 and 9, along with the normalized standard deviations in the table below to quantitatively measure the variations of errors. The results show that TabDiff with learnable schedules not only achieves lower average error but also exhibits more consistent errors across columns. This balance indicates that learnable schedules help the model balance its capacity on different dimensions of the data, improving the model’s ability to handle heterogeneous distributions.

Thank you for your constructive feedback and insightful questions! Below, we provide detailed responses to each of your questions. Additionally, please refer to our revised paper, with updates colored in purple.

[W1, W2] Detailed illustration for TabDiff’s diffusion processes and the modeling of categorical columns

Thank you for pointing this out. We completely agree that the presentation of our model’s training and sampling process would benefit from an intuitive illustration.

Now in Appendix C (“Detailed Illustrations of Training and Sampling Processes”) of the revised manuscript, we closely illustrate the diffusion training and testing processes using an intuitive visualization in Figure 5 that showcases the process of generating a toy example. We also more closely described the modeling of categorical columns by referring to specific examples in this illustration.

[W3] Missing parenthesis in Eq.(10)

Thank you for pointing this out. We have corrected it in the revised manuscript!

Dear Reviewer zTT5,

We would like to express our gratitude for your valuable comments and suggestions. Thanks a lot for your positive feedback on our method and experiments.

As the author-reviewer discussion period is coming to a close on Nov 27th, we would greatly appreciate it if you could clarify any remaining concerns you may have. This will ensure that we can adequately address them during this period.

Thank you very much for your attention to this matter.

Best regards,

Authors

Thank you for your constructive feedback and insightful questions! Below, we provide detailed responses to each of your questions. Additionally, please refer to our revised paper, with updates colored in purple.

[W1, Q1] Cross-modal modeling

Thank you for raising the point regarding how TabDiff models numerical and categorical features together. We elaborate on this point in the following.

Are the features modeled separately? First, we would like to clarify that, TabDiff does NOT model different types of features separately, though it uses different diffusion processes for them. TabDiff models the joint of all features under the same continuous time diffusion framework, where all features are perturbed with noise together in the forward training process and denoised together in the reverse sampling process. The diffusion process is designed differently in order to accommodate the distinct natures of each data type (e.g. we destroy the discrete nature of categorical features if we add gaussian noise to them).

How modelings of different features help each other. The modelings of numerical and categorical features help each other by capturing their mutual dependencies within the joint distribution. Such dependencies are captured in TabDiff in two aspects.

First, in the reverse diffusion process, we adopt the common practice in multimodal diffusion models by factorizing the posterior across different features, as described in Eq.(2). This factorization does not imply that the features are independent. Each conditional distribution $q(x_s^{\text{num}} \mid x_t, x_0)$ and $q(x_s^{\text{cat}} \mid x_t, x_0)$ is conditioned on the entire $x_t$ and $x_0$, which contains information from both numerical and categorical features. This allows the model to learn and capture complex dependencies between different features.

Beyond this standard practice, the modeling of dependencies between different features is enhanced by the feature-wise adaptive noise schedules. These learnable schedules give the model the flexibility to denoise different features at varying speeds. Specifically, a more positive $\rho$ and $k$ corresponds to quicker denoising. During training, the model learns different denoising rates for different features feature, which can be regarded as an order of generation. Features with larger $\rho$ and $k$ are recovered earlier in the denoising process, allowing subsequent features to be generated and conditioned on them. This mechanism enables the model to better capture inter–feature dependencies. In Table 2, TabDiff outperforms all other baselines on inter-column correlation metrics (i.e., Trend); in Table 5, TabDiff with learnable schedules achieves higher scores on Trend.

Dear Reviewer gKXC,

We would like to express our gratitude for your valuable comments and suggestions. Thanks a lot for your positive feedback on our method and experiments.

As the author-reviewer discussion period is coming to a close on Nov 27th, we would greatly appreciate it if you could clarify any remaining concerns you may have. This will ensure that we can adequately address them during this period.

Thank you very much for your attention to this matter.

Best regards,

Authors

Hi Markus,

Thank you for bringing the related work to our attention. We are delighted to discuss our specific modeling choices and believe that this discussion will deepen understanding for both you and us!

[Q1] Regarding the balancing impact of feature-specific noise schedules

Yes, we do observe such an impact. As detailed in Appendix F.2, TabDiff with feature-specific schedules demonstrates more uniform column-wise error distributions compared to baseline approaches.

Regarding how the different scales of continuous and discrete loss impact this balancing effect: this is a very good consideration. The different scales in loss could indeed have an effect. But, in our experiments, we haven’t observed how unbalanced loss impacts the balancing effect. This could be attributed to the fact that we thoroughly trained both modalities thoroughly with designs such as the loss calibration techniques as discussed in Q3. However, we do consider that a more thorough study of loss scales’ impact would be valuable for future work, particularly in extreme cases where the loss scales differ by several orders of magnitude.

[Q2] The case when a categorical feature depends on the fine-grained details of some continuous features

This is an insightful question. First, we want to clarify that TabDiff’s sampling process, as outlined in Algorithm 2, does allow unmasked categorical features to potentially change their values in later sampling steps. In line 6 of Algorithm 2, our stochastic sampler first stochastically perturbs the current sample to a slightly higher noise level, before making the prediction. Thus, an unmasked categorical feature might become masked again and get re-predicted at a later timestep, when the fine-grained details of the continuous features are better built.

Even if the stochastic sampler is disabled, we haven't observed performance degradation in such cases in our experiment. This might be attributed to the fact that feature dependencies are typically flexible and arbitrary, with no specific order. If categorical features depend on the fine-graded details of continuous features, then the reverse is also true. Given this, our generation process can accommodate such a case: when a categorical feature is unmasked early in the generation process, the continuous features and other unmasked categorical features would need to adjust their own content based on the fine-grained details provided by the unmasked feature.

[Q3] Decaying continuous feature loss weight

Yes, your rationale is exactly correct. In our experiment, we observed that under our current loss scale, the discrete loss takes a bit longer to get thoroughly trained than the continuous. We speculate this opposite phenomenon, relative to your finding, might be due to the difference in the how model's categorical features or the difference in our loss scales.

Given our observation that the training loss has a direct correlation with sample quality, we would want a technique that thoroughly trained the discrete feature while avoiding overfitting the continuous features. A linearly decreasing weight loss turns out to be an effective and robust solution that generalizes well across all datasets.

We hope this addresses your questions!

Authors

Dear AC,

Thank you for your feedback and we have addressed your two requirements accordingly for our camera-ready version. Please let us know if there are other concerns!

Sincerely, Authors