KBFormer: A Transformer-based Diffusion Model of Structured Entities with Heterogeneous Properties

Thank you for the thorough review and suggestions!

1. Although the authors make clear introductions to the hybrid diffusion training paradigm, the explanation for the model architecture is not clear enough (sometimes even confusing). In fact, the first of my two questions is because of not understanding the architecture here. I suggest the authors can modify Figure 5 and 6…

We have made significant strides in improving the clarity of the text and the model architecture in particular (see changes in blue). Following your suggestions, we completely reworked Figures 5 and 6 (now combined into Figure 4).

2. Although the inherent probabilistic nature of the model makes it suitable for prediction tasks for scientific scenario, I still believe that comparisons between the KBFormer and some regression model baselines should be included to demonstrate the effectiveness of the model. Besides, in order to illustrate the model's capability in scientific applications, more experiments on benchmarks from other disciplines like molecules, proteins, etc., are required.

Thank you for the suggestions! We have now included an additional baseline: Gradient Boosted Decision Trees (GBDT)s, which offer state-of-the-art performance on tabular data for both regression and classification tasks. Our experiments show that we can perform competitively against specialized regression models even though our model is generative in nature. We also included additional experiments across 15 datasets and compared our results against other generative models spanning various tasks, including two scientific datasets (MiniBooNE and Higgs).

Q1. In terms of the model architecture, … why it is necessary to use a transformer model with such a short "sequence length".

We have updated this figure and added extensive details around it. The previous figure used a small example with three properties for simplicity, though in practice, we use many more (see Table 6 with a description of all the datasets we used and the number of properties in each one.) The updated figure should clarify that there are more properties than what is shown. The transformer aggregates information across all properties (e.g., our tabular data experiments have up to 50). In future work, we aim to scale to a much larger number of properties and knowledge bases, and the transformer architecture is a good candidate for this.

Q2. Transformer encoders are used to encode text fields. However, if the text is simply a property value (e.g. "iPhone" in Figure 5), why not just use a pretrained word embedding?

Figure 5 (now 4) should no longer imply that text fields are single words. If this were the case, we could indeed use pre-trained word embeddings. However, the goal of the text generator is to fill text property values with a (possibly long) sequence of tokens.

Q3. The diffusion training strategy is a promising solution for probabilistic-based generative models. However, is the diffusion paradigm the optimal method to do this? Is there any other generative model (or even regressive model) that does not rely on the noise and denoise process?

In our latest round of experiments, we explored the performance of our diffusion-based approach against several other generative models including VAEs, GANs, and a different implementation of diffusion on tabular data (TabDDPM). We also compare against regression models for nuclear physics and the GSM phone dataset. Across all the experiments, KBFormer obtains favorable performance, which illustrates the strength of the approach. Please note that our diffusion process has “less noise” than expected from more standard approaches. When we make predictions autoregressively, noise is injected via the random order in which properties are unmasked.

Thank you for your review! We uploaded a revised paper with additional clarifications, baselines, and experiments, which we hope address your concerns.

1. The paper's contribution seems to be a bit incremental, since the diffusion modeling over heterogeneous data in Section 3.2 follows the previous work [1]. It would be helpful if the authors could clarify the difference/contribution of the proposed method.

Here we summarize the differences between our work and [1]:

D3PM [1] deals with discrete-state spaces in discrete time. Campbell’s formulation [2] (which itself is parallel to [1]) gives a very general case for continuous-time discrete-state models. Our derivation builds on the continuous-time approach and shows that for the absorbing-state case, the ELBO can be estimated using a single evaluation of the model (two are required in Campbell’s formulation, or a biased approximation). Additionally, [1] and [2] deal with continuous state spaces via Gaussian discretization, whereas we propose using a Gaussian Mixture Model.

Moreover, in addition to the diffusion paradigm we derive, we also propose an architecture and hierarchical encoding scheme that can handle structured data and obtain superior performance across most datasets. Finally, note that in many of our new experiments, we compare against TabDDPM, which is partly an implementation of [1] to handle tabular datasets, and show that we obtain favorable performance in many cases.

In the experiments, only the baseline that always predicts the marginal mode/mean is compared in terms of the prediction accuracy with different unmasked rates. Are there any other baselines with regression/autoregression that can be compared? Also it would be helpful if the authors could elaborate on Section 3.1 about why other traditional models with regression and masked modeling will have less optimal performance.

Yes! We added Gradient Boosted Decision Trees (GBDTs) as a new baseline in our reconstruction experiments. We train a GBDT on all other properties for each numerical and categorical type. This is a much stronger baseline than what we had before (see the updated section 5 and appendices E and F, which contain more training details). We also include experiments on 15 datasets with new baselines built on various generative frameworks (GAN, VAEs, DDPM, etc).

As for section 3.1, we argued that generating many properties at once can degrade performance. In our diffusion framework, that would be akin to skipping “time steps.” We ran ablations on the 15 datasets to show that generating all properties at once can impact performance greatly (see Table 4 in Appendix B.1). We also show a few qualitative examples in Appendix B.1.2). In one of these examples, we generate MNIST digits in one time step, and one can visually inspect the samples to see they are a lot noisier compared to the autoregressive samples. To clarify further, suppose in a language modeling setting, we are interested in predicting a sequence of two tokens that answers a question positively. “[Of] [course]” and “[No] [problem]” could be equally likely candidates. However, sampling the two tokens simultaneously can result in answers like “[No] [course]” or “[Of] [problem].” We run exactly into this inconsistency issue if we try to sample all properties simultaneously. But it can be solved by conditioning the second token generated on the first token. This is why our autoregressive model performs much better than the single-step approach.

[1] Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34:17981–17993, 2021.

[2] Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 35, 2022.

Thank you for your review and all the suggestions! We are currently working on implementing them. We would be happy to revise our references with additional material. Are there any specific references we are missing?

Thank you for your review and the great suggestions. The new version of the paper reflects the changes mentioned below.

• This paper lacks ablation studies. The paper does not include ablation studies to analyze the contribution of different components of the model to its overall performance. For example, it is mentioned in paragraph "Encoding" of section 4, that two alternatives for embedding numerical values, yet it lacks a quantitative performance comparison between these methods. Conducting such an analysis could shed light on the critical components of the model and direct future enhancements.

You are absolutely right! We have added ablations in Appendix B for the following:

• Single-step masked modeling vs. absorbing state (autoregressive) diffusion.
• DICE vs. Periodic Encoding for the numerical embeddings.
• GMM vs. MSE for numerical quantities. (Using GMM with 1 mixture and unit variance is equivalent to training with MSE.)
• Tying numerical embeddings across properties.

• There is no discussion in the paper about the limitations or potential failure modes of the KBFormer method. Including this could be crucial for fully understanding where the model may fall short in practical applications and where further research and development could be most beneficial.

When mentioning future work, we allude to scaling to larger datasets and performing pre-training across multiple domains. That is our current main limitation. We made it more explicit in the paper and added a longer limitations section in Appendix G.

• I would like to suggest the authors confirm whether the abbreviations used throughout the paper, such as “KB” in the abstract, are consistently defined upon first use? A uniform approach to abbreviation would aid reader comprehension.

This should be resolved in our latest update of the paper. Apologies for missing some definitions in the first version!

• The introduction promises that the KBFormer model addresses several tasks: KB completion, entity linking, anomalous property detection, and enhancement of foundation models with learned representations. The experiments, however, seem to showcase a subset of these. Can the authors clarify the criteria for experiment selection and indicate if demonstrating the model's capabilities on the remaining tasks is within the scope of future work?

In the introduction, we list some applications to motivate developing our approach, though it is true we do not tackle all tasks we use for motivation. However, we are working on extensions to some of the other tasks in future work.

• Would the authors consider revising the reference list to ensure that all citations are consistent and reflect the most current research where applicable? This would help maintain the paper’s relevance and assist readers in locating the sources.

We updated our references for improved consistency across references from the same source.

Cheers! The authors

Thank you for your review! We elaborated extensively in the paper (see changes in blue). We hope you find the updated Figure 5 (now Figure 4) to be clearer. If not, please let us know what points remain unclear so we can clarify them both in the paper and our final response.

We appreciate your thorough feedback!

1. The paper does not discuss or compare with other methods that can handle discrete …

Thank you for bringing this literature to our attention. In this submission, we focus on handling different property types in a generative framework; in future work, we will focus on knowledge-intensive tasks, making these comparisons very important.

In the latest submission, we conducted additional experiments across 15 datasets, comparing KBFormer against several other generative models specializing in structured data. The additional results showing favorable performance on most datasets against CTABGAN, TVAE, and TabDDPM are presented in Table 3. We hope that these new experiments satisfy your request for more baselines.

2. The section on “Continuous Relaxation of Discrete State Diffusion” is not well explained…

We have changed the contents of this section and added significantly more details to improve the clarity and the general flow (see the new text in blue).

To summarize here, the objective of this section is to give some intuition on how one might utilize the formulation of discrete space diffusion for continuous properties. Naively, if one wants to model a continuous property as a multi-class prediction problem, one would discretize the space with many bins (to maintain good precision), but the softmax computation can become quite expensive. This would also discard ordinality. However, under the simplifying assumption of Gaussian properties, we can formulate the problem as a mean-squared error (MSE) and maintain ordinality and good precision. This assumption does not usually hold, so we must use something more expressive, like a Gaussian Mixture Model (GMM), which can be seen as a generalization of the MSE approach. This section aims to develop intuition to justify using GMMs to model numerical properties within the full framework.

In practice, we completely drop any discretization and use the GMM parameters to sample numerical properties instead of logits. The number of mixtures we choose is simply a hyperparameter.

Please let us know if anything remains unclear so we can elaborate in further responses and the paper.

3. The paper needs to improve its writing quality and clarity. Some specific issues are: Proposition 1: …

We expanded our explanations and improved the general flow (see changes in blue). We now clearly point to the proof of Proposition 1 in the appendix, and we fixed all the typos/issues we found.

Q1. Which type of encoder and decoder did you use for different property types: i)- conditioning on the property itself or ii) disjoint encoders for each property?

One can do both. On a larger scale, we expect conditioning to enable dealing with unseen properties, but for the experiments in this paper, we used disjoint encoders/decoders. We ablate performance when sharing numerical embeddings in the current framing on GSM, which also seems to work.

Q2. In page 6, it is stated that “ ‘year’ has the same representation in different contexts …

Yes, that was unclear. Apologies for the confusion. The GSM dataset has the property "Phone.Launch.Year". The RNN will process this as a sequence, RNN(["Phone", "Launch", "Year"]). Thus, if the token "Year" exists in another context (let's say "Phone.Discontinued.Year"), the RNN representations will contain information on the similarity of both keys.

Q3. For experiments in section 5.2, what is the evaluation method? It is also worth…

We added a strong baseline across both sections referenced. In (prior) section 5.2, we evaluated the model's predictions on a given property when conditioned on all other available properties (this dataset is sparse).

In (prior) 5.1, we explained why we rotate the inputs as an augmentation. To answer the question here, since LLaMA is sequence-autoregressive, we have to rotate the string representation of the dictionary such that each field/property can appear in the last position. This ensures we can predict a value for the property while conditioning on all other properties.

More generally, causal models can perform worse when predicting properties [1], even when they are in the training set. So, we ensure that the model has seen enough permutations of the properties to be able to make predictions while conditioning on any subset.

[1] Physics of Language Models: Part 3.2, Knowledge Manipulation