Paper: KBFormer: A Transformer-based Diffusion Model of Structured Entities with Heterogeneous Properties (ICLR-2024)

Reviewer comment (Experiments):
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.

Author response:
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).

Reviewer comment (Experiments):
In terms of the model architecture, I suppose the output from the encoder is $(3, feature size)$ for the example shown in Figure 5. If that is true, I'm wondering why it is necessary to use a transformer model with such a short "sequence length".

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Novelty):
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?

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Writing):
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.

Author response:
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?

Reviewer comment (Presentation):
The presentation of the paper is quite poor. For example, Figure 5 is not comprehensive, with many parts unexplained.

Author response:
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.

Reviewer comment (Writing):
The section on “Continuous Relaxation of Discrete State Diffusion” is not well explained. It is unclear what its objective is; is the goal to learn bin centers, and how they are used in demasking? Is the discretization with 256 bins and learned bin centers similar to GMM with 256 mixtures? The paper also introduces some terms without proper definitions, such as “an infinite bin limit approximation” and “discretization with a large but finite bin density”. It would be helpful to provide more details and intuition behind these concepts.

Author response:
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.