TabNAT: A Continuous-Discrete Joint Generative Framework for Tabular Data

We sincerely thank the reviewer for the thoughtful and constructive feedback on our manuscript. We appreciate the time and effort invested in providing these valuable comments, which will help improve the quality of our paper. Below, we address the points raised:

Choice of Experimental Datasets

The reviewer raised concerns about our dataset selection, particularly noting differences from those used in DP-TBART/Tab-MT papers. We would like to clarify that our experimental design primarily followed TabSyn's pipeline, as this was our initial baseline for comparison. Accordingly, we incorporated the six heterogeneous datasets used in the TabSyn paper to enable direct performance benchmarking.

To demonstrate TabNAT's versatility across different data types, we deliberately expanded our evaluation by adding two datasets containing only continuous columns and two datasets with only discrete columns. This comprehensive approach allows us to assess TabNAT's generalizability across diverse tabular data structures.

We acknowledge the reviewer's valid point regarding datasets used in TabMT and DP-Tbart. We have conducted additional experiments using five additional datasets used in these papers, with results presented in rebuttal.pdf in the anonymous link: https://anonymous.4open.science/r/ICML-TabNAT. These supplementary experiments further validate TabNAT's effectiveness across a broader range of benchmark datasets.

Performance Comparison with TabSyn

We appreciate the reviewer's insightful observation regarding TabNAT's performance relative to TabSyn across different datasets. Our detailed analysis reveals several important patterns in performance differences:

First, TabNAT consistently outperforms TabSyn on both continuous-only and discrete-only datasets. This advantage stems from TabNAT's architecture, which employs specialized generation models optimally suited for each data type. Unlike TabSyn's approach of using a VAE encoder to create a latent space, TabNAT processes different feature types more directly and naturally, minimizing information loss that typically occurs during VAE encoding.

For mixed-type datasets, our analysis reveals that TabNAT demonstrates clear advantages when the proportion of discrete features is higher, as evidenced by superior performance on the Adult and Beijing datasets. This performance differential can be attributed to TabNAT's transformer-based categorical generation component, which better captures the complex dependencies between categorical variables through its self-attention mechanism. Additionally, TabNAT's architecture enables more effective modeling of interactions between continuous and discrete features, particularly when discrete features play a dominant role in the underlying data distribution.

Importantly, even in datasets where TabNAT and TabSyn demonstrate comparable statistical performance, TabNAT offers significant practical advantages that TabSyn cannot match. Specifically, TabSyn's VAE-based architecture fundamentally limits its ability to perform flexible conditional generation. In contrast, TabNAT's arbitrary-order autoregressive generation capability supports any form of conditional generation, including missing data imputation. This flexibility stems from TabNAT's bidirectional masking strategy and specialized components for different data types, allowing users to specify any subset of features as conditions while generating the remaining features—a crucial capability for many real-world applications.

Theoretical Analysis

We thank the reviewer for suggesting a theoretical analysis to support TabNAT's empirical performance. Our approach is grounded in fundamental probability theory, specifically the decomposition of joint distributions as expressed in Equations (1) and (2). TabNAT's key theoretical contribution lies in transforming the tabular data generation problem into a series of conditional distribution modeling tasks through our bidirectional Masked Transformer architecture.

This decomposition allows us to leverage the most appropriate generative modeling techniques for different variable types: diffusion models for continuous variables and autoregressive models for discrete variables. Both approaches have strong theoretical foundations in distribution modeling. Diffusion models offer provable convergence to the target distribution through a gradual denoising process, while autoregressive

Minor Corrections

We have corrected the noted typographical errors:

• "After TabNAT is well-trained" → "After TabNAT is well trained"
• Removed the extra whitespace in "TabNAT 's training speed"

Additionally, we have thoroughly reviewed the manuscript to ensure consistency in terminology and formatting throughout.

We sincerely thank the reviewer for their thoughtful assessment of our manuscript. We appreciate the recognition of our paper's strengths, including its clear writing, promising results, and well-supported claims regarding statistical fidelity, utility, and privacy preservation. Below, we address each of the reviewer's questions and concerns.

Response to Comments on Figures

We thank the reviewer for the suggestion regarding Figures 1 and 2. We agree that Figure 3 provides better context and will consider restructuring the presentation of figures in the revised manuscript to improve clarity.

Position-specific mask embedding and position embeddings

Intuitively, the position-specific mask embedding serves a different purpose than the position embeddings. While position embeddings encode the absolute position of each token in the sequence, the position-specific mask embedding is designed to capture information about which specific attributes are masked during the any-order training process. This allows the model to better understand the relationship between masked and unmasked attributes in different positions, enhancing its ability to handle partial information scenarios. In fact, we found that the empirical performance difference between the two designs is not substantial (with position-specific mask embedding performing slightly better). Considering that the additional parameters introduced by this design are negligible, we adopted this approach in our final model.

MLE task performance clarification

Thank you for raising this important point. We apologize for not including these results in Table 1. The complete results for classifiers and regressors trained on the original data can be found in the Appendix, page 20, Table 14. We commit to revising Table 1 to include this information in the main text for better clarity.

Regarding the number of synthetic samples used, the results reported in the last two columns of Table 1 were obtained using the same number of synthetic examples as in the training set, which varies according to the size of each dataset. This 1:1 ratio ensures a fair comparison across different methods. This information is provided in Appendix C.7.5.

DCR metric definition

We sincerely apologize for only providing the experimental setup for DCR without explaining how this metric is calculated.

DCR (Distance to Closest Record) is a commonly used metric for evaluating the privacy protection performance of synthetic data. For each synthetic example, we calculate its minimum distance to any example in the training set.

Formally, let's assume we have $N$ training examples {$r_1$, $r_2$, ..., $r_N$} and $n$ synthetic examples {$s_1$, $s_2$, ..., $s_n$}. For each synthetic example $s_i$, we compute:

$\text{DCR}^{train}(s_i) = \min_{j \in \{1,2,...,N\}} d(s_i, r_j)$

where $d(\cdot,\cdot)$ represents the normalized Euclidean distance between two records after appropriate feature normalization.

Similarly, for $N$ holdout examples {$h_1$, $h_2$, ..., $h_N$}, we can compute their DCRs:

$\text{DCR}^{holdout}(s_i) = \min_{j \in \{1,2,...,N\}} d(s_i, h_j)$

Since the training set and holdout set are i.i.d. samples from the same distribution, if the synthetic samples learn the underlying distribution from the training set correctly, $\text{DCR}^{train}$ and $\text{DCR}^{holdout}$ should be similar. Otherwise, if the synthetic examples are copied from the training data, the $\text{DCR}^{train}$ of each synthetic example $s_i$ should be close to zero.

Setting of missing data imputation

Thank you for raising this point about our imputation task setup. We would like to clarify that in real-world scenarios, we often face in-sample missing data imputation tasks, such as during the data preprocessing phase of large data science projects.

Both in-sample and out-of-sample imputation approaches have been explored in previous research [1,2]. Methods like HyperImpute adopt out-of-sample evaluation primarily because, as a hybrid approach, it requires supervised labels for model selection. Our method does not have this requirement. We also note that works like Remasker seem to adopt the in-sample setting as well.

That said, we agree that out-of-sample imputation is equally important. Our model is naturally suitable for out-of-sample settings, and we will conduct additional experiments to demonstrate its effectiveness in this context. We have included these results in rebuttal.pdf in the anonymous link：https://anonymous.4open.science/r/ICML-TabNAT to provide a more comprehensive evaluation of our method's imputation capabilities across both settings.

We sincerely thank the reviewer for the thorough review and constructive feedback on our manuscript. We have addressed the raised concerns as follows:

The suitability of diffusion modeling for continuous columns in tabular data

We agree that this is an important consideration that deserves clarification. Our choice of diffusion modeling for continuous columns is supported by empirical evidence from recent SOTA work:

• TabDDPM was the first model to successfully apply diffusion models to tabular data synthesis, using traditional DDPM for continuous columns and discrete diffusion with multinomial distribution for discrete columns. Its strong performance validated the effectiveness of diffusion models for capturing the complex distributions in continuous (multi-column) tabular data.
• TabSyn further advanced this approach by mapping both discrete and continuous columns to a continuous embedding space, enabling the application of standard diffusion models to the entire table. TabSyn has achieved state-of-the-art results in tabular data generation, demonstrating the exceptional capability of diffusion models to capture the distributions of continuous tabular data.
• Our empirical experiments align with these findings. It's important to note that direct autoregressive generative modeling of continuous values is not feasible without discretization. Baseline methods such as Tab-MT and DP-TBART both employ this discretization-plus-autoregressive approach, and our experiments demonstrate that they consistently underperform compared to our method across all metrics, particularly in Statistical Fidelity scores.

On the differences between continuous columns in tabular data and images

We agree with the reviewer that a deeper analysis of the differences between continuous values in tabular data and images would strengthen our paper. These differences are fundamental to our methodological choices:

• Tabular data columns are fundamentally more independent than image pixels, with each column typically representing a distinct feature with its own semantic meaning. In contrast, image data pixels have strong spatial continuity, with neighboring pixels highly correlated and collectively representing coherent visual elements.
• This structural difference necessitates different denoising neural network architectures. For images, CNN-based UNet models have become the standard architecture for diffusion models because they effectively leverage the local spatial correlations through convolutional operations and capture multi-scale features through their encoder-decoder structure with skip connections.
• For tabular data, where columns have distinct meanings without inherent spatial relationships, traditional MLP architectures are more appropriate as denoising neural networks. In our approach, we use an MLP for denoising continuous data, while the Transformer is employed solely to generate conditional vectors. This design choice aligns with other successful tabular diffusion models, such as TabDDPM and TabSyn, which also utilize MLPs as their denoising neural networks.

On the missing data imputation experiment

We would like to clarify that the baselines we compared against (Remasker, HyperImpute, and GRAPE) are indeed SOTA methods for tabular missing data imputation. These same methods were used as the top-performing baselines in a recent ICLR 2025 Spotlight paper on tabular data imputation [1]. Our selection of baselines is therefore aligned with the current standards in the field. We will aslo include the new paper as an additional baseline in our revised manuscript to ensure our comparisons remain comprehensive and up-to-date.

Due to page limitations, we were unable to include ablation studies for the missing data imputation task in the original submission. We have now conducted these additional experiments, and the results are presented in Figure 15 and Figure 16 in rebuttal.pdf in the anonymous link：https://anonymous.4open.science/r/ICML-TabNAT. The findings are summarized as follows:

• Simple MSE loss is as effective as the Diffusion loss in missing data imputation since we care more about the expectation of the missing entry rather than its distribution. The imputation of discrete columns is sensitive to the order, and the random-order sampling used in our paper is beneficial.

• The proposed TabNAT is very robust to the model depth and width in the missing data imputation task.

These results further validate our design choices and demonstrate that each component of our model contributes significantly to its superior performance in the missing data imputation task.

Missing references

We thank the reviewer for pointing out the related works. We will cite these works and add a discussion about the application of diffusion losses in autoregressive image modeling.

References:

[1] Zhang et al. DiffPuter: Empowering Diffusion Models for Missing Data Imputation. In ICLR 2025.

We sincerely thank the reviewer for their thoughtful and constructive feedback on our manuscript. We appreciate the time and effort spent reviewing our work, and we believe the suggestions will significantly improve the quality of our paper. Below, we address each point raised by the reviewer.

Modeling of Continuous Data

Your observation about our approach to modeling continuous data is insightful. We agree with your assessment that our method doesn't directly model the full continuous distribution $p(X)$, but rather learns a conditional distribution $p(X|c)$, where the constant vector $c$ represents the case where all discrete columns are masked.

While this is conceptually different from modeling the complete joint distribution, in practice, the two approaches yield equivalent results for generation purposes. This is because when all discrete columns are masked, the model effectively sees only a generic placeholder with no specific information, forcing it to generate continuous values based solely on the learned marginal distribution of those values. The mask tokens in this case serve merely as positional indicators without conveying any actual discrete data information.

Your suggested approach of encoding perturbed continuous dimensions into vectors at positions 1 to K is elegant and theoretically sound. Following your suggestion, we implemented and tested this approach. Our experiments confirmed your intuition—the results were indeed comparable to our original method in terms of generation quality.

However, we identified a practical limitation: this approach significantly extends the sequence length, approximately tripling the training computational cost. Given the similar performance outcomes but increased resource requirements, we opted for our original approach as a more efficient solution. Nevertheless, we appreciate this valuable suggestion and will discuss this alternative formulation and its trade-offs in our revised paper.

Discrete Diffusion Models

Thank you for highlighting the recent advances in discrete diffusion models that we overlooked. While we did include basic approaches like multinomial diffusion in our baselines (which showed limited effectiveness), we missed the opportunity to explore more recent developments in this area.

The papers you referenced indeed offer promising directions for modeling discrete data without left-to-right bias. In our revised paper, we'll discuss these advancements and explore how they might complement our approach for mixed-type tabular data generation. We're particularly intrigued by masked diffusion approaches and their potential in heterogenous tabular data generation.

Choice of Stochastic Generation Approach

We thank the reviewer for this insightful suggestion. Our primary focus was to effectively integrate diffusion models with autoregressive approaches for tabular data generation. Regarding the choice of sampling method, we acknowledge that different approaches offer various trade-offs between quality and efficiency. Following the reviewer's suggestion, we conducted additional experiments with flow matching. While this approach did improve sampling speed (requiring only 20 steps), we observed a slight decrease in generation quality. It's worth noting that our current sampling method requires only 50 sampling steps—significantly fewer than the 1000 steps typically needed in traditional DDPM approaches—thus already providing a good balance between quality and efficiency.

Position-Specific Masking Encoding

Intuitively, the position-specific mask embedding serves a different purpose than the position embeddings. While position embeddings encode the absolute position of each token in the sequence, the position-specific mask embedding is designed to capture information about which specific attributes are masked during the any-order training process. This allows the model to better understand the relationship between masked and unmasked attributes in different positions, enhancing its ability to handle partial information scenarios. In fact, we found that the empirical performance difference between the two designs is not substantial (with position-specific mask embedding performing slightly better). Considering that the additional parameters introduced by this design are negligible, we adopted this approach in our final model.

Notation and Figure Clarity

We thank the reviewer for pointing out the notation issues and concerns about figure clarity. We will correct all notation errors and improve the clarity of Figures 1 and 2 in the revised version of our manuscript.