Towards Foundation Models for Knowledge Graph Reasoning

We thank the reviewer for highlighting the proposed method, impressive experimental results, and the quality of writing. Below, we would like to comment on the raised weaknesses and questions.

W1. The proposed method relies entirely on knowledge graph structure and does not consider using node embedding such as textual features of the knowledge graphs. In reality, text embedding of nodes and edges could be a better transferrable embedding.

It is worth noting that LLM-processed text features are only available for KGs that do have those entity/relation descriptions available, but not all KGs have such text features readily available. In the case of numerical features, LLMs might be even a suboptimal encoder choice, eg, encoding numbers might require a different encoder, and if numerical features adhere to certain symmetries, then geometric models might be a better choice. That’s why in this work we focus on the structural representations and feature-less graphs as this can be applied to any KG with or without features. We briefly cover those reasons as well as some recent text-based approaches in the Related Work -> Text-based methods subsection.

Having said that, there is some evidence (like in ReFactorGNNs [1]) that concatenating encoded text features (where available) to structural GNN features is likely to further boost the performance in inductive tasks. This is an intriguing direction for the future work and we elaborated on that in the new Appendix E in the updated version.

Thank you for mentioning PRODIGY, it is certainly a related work (although focusing on different tasks of node classification and relation prediction) and we added it in the updated version of the manuscript.

W2. The model does not scale well as the authors already pointed out.

We acknowledge the scaling behavior in the limitations and consider it a very promising avenue for future work. In particular, scaling laws for GNNs and common graph learning tasks (like link prediction) are not derived yet so we can only hypothesize whether there is any connection between GNNs size, dataset size, graph topology, and expected performance. Generally, there is no consensus in the graph learning community on whether deep or wide (non-geometric) GNNs bring immediate benefits - mostly due to the rising issues of oversmoothing and oversquashing (some initial results were recently presented in [2]). In our experiments, we observe that the diversity of graphs in the pre-training mixture plays an important role as well. Therefore, we believe that a brute-force increase of the model size is unlikely to bring benefits unless paired with more diverse training data and more intricate mechanisms for capturing relational interactions.

W3. The zero shot and fine-tuning performances are worse or on-par with the per dataset model performance, rendering pretraining not effective performance-wise.

We would like to offer a different perspective on this matter (perhaps even the opposite) - if a single pre-trained model can deliver a comparable performance to 57 models trained individually from scratch on each dataset, then it saves a significant amount of compute needed to train those individual models. Zero-shot performance of ULTRA highlights such computational efficiency whereas fine-tuning requires only a fraction of the full training costs, as reported in Table 8 in the Appendix. Generally, we consider the ability of a single model to match the performance of numerous dataset-specific models to be promising. As we are at the beginning of the transfer learning era for KG reasoning, we believe that more fine-grained pre-training and fine-tuning techniques might further improve the compute/performance trade-off and it is an exciting area for the future work.

Q1. What are $u$ and $v$ in $h_{u|v}$ in section 4.2?

Thanks for noticing that. Indeed, in Section 4.2 the notation of the indicator function and message passing can be simplified (we changed that in the updated version of the manuscript). In Section 4.2 which deals with the relations graph, the indicator function says that a node $v$ in the graph of relations will be initialized with some non-zero vector (vector of all ones in our case) if it’s equivalent to the query relation $q$ in the original input query $(h,q,?)$. The notation for a node $u$ is indeed redundant in this case as we can condition both equations on the query relation $q$ directly.

We thank the reviewer for highlighting the merits of our work as to the importance of the problem, the ability to perform zero-shot inference and fine-tuning, as well as the quality of writing.

Please find our comments on the weaknesses and questions below.

W1. Could you elaborate more on the process and result of the indicator function according to those two variables, perhaps with visuals?

Indeed, in Section 4.2 the notation of the indicator function can be simplified (we changed that in the updated version of the manuscript). In Section 4.2 which deals with the graph of relations, the indicator function says that a node $v$ in the graph of relations will be initialized with some non-zero vector (vector of all ones in our case) if it’s equivalent to the query relation $q$ in the original input query $(h,q,?)$. The notation for a node $u$ is indeed redundant in this case as we can condition both equations on the query relation $q$ directly.

W2. I think one potential direction … might be to use the LMs, and the authors may highlight this point more and potentially make comparisons between the proposed approach and text-based methods. I don't think this should be the critical weakness of this paper since text-based methods are limited to knowledge graphs with textual features;

You correctly pointed out that LLM-processed features are only available for KGs that do have those entity/relation features available, we briefly cover some recent approaches in the Related Work -> Text-based methods subsection. Not all KGs have such text features readily available and in the case of numerical features LLMs might be even a suboptimal encoder choice. That’s why in this work we focus on the structural representations and feature-less graphs as this can be applied to any KG with or without features.

Having said that, there is some evidence (like in ReFactorGNNs [1]) that concatenating encoded text features (where available) to structural GNN features is likely to further boost the performance in inductive tasks. We consider dataset-specific features complementary to ULTRA representations and hypothesize that such additional features might be particularly useful at the fine-tuning stages. This is an intriguing direction for the future work and we added this discussion in the new Appendix E in the updated version.

Q1. Previous methods are superior to ULTRA on transductive graphs, which are worthwhile to discuss.

Thank you for the suggestion, we added more discussion on the performance on larger transductive graphs in Appendix D. Generally, we attribute the performance difference to the following factors:

• Training data mixture and OOD generalization: the model reported in Table 1 was trained on 3 medium-sized KGs (15k - 40k nodes, 80k - 270k edges) while the biggest gaps are on larger graphs with many more nodes and edges (up to 120k / 1M for YAGO 310), or many more relation types (1600+ in AristoV4), or very sparse (as in ConceptNet100k with 100k edges over 78k nodes). Size generalization issues are common for GNNs as found in [2] and [3]. However, if we take the ULTRA checkpoint pre-trained on 8 graphs (Table 9 in Appendix D) and run evaluation on all 16 transductive graphs, then the average performance is better than supervised SOTA models.

Of course, for those datasets in the pretraining mix, we cannot call this regime a zero-shot inference anymore, but the results hint upon the hypothesis that more diverse graphs in the training mixture are beneficial for alleviating OOD issues.

• Transductive models have the privilege of memorizing target data distributions (even on large graphs) into entity/relation-specific vectors with overall many millions of parameters, eg, 80M parameters for a supervised SOTA BiQUE on ConceptNet100k. This performance, however, comes with the absence of transferability across KGs. In contrast, all pre-trained ULTRA checkpoints are rather small (~170k parameters) but generalize to any KG. When it comes to scaling, scaling laws for GNNs and common graph learning tasks (like link prediction) are not derived yet so we can only hypothesize whether there is any connection between GNNs size, dataset size, graph topology, and expected performance. Generally, there is no consensus in the graph learning community on whether deep or wide (non-geometric) GNNs bring immediate benefits - mostly due to the rising issues of oversmoothing and oversquashing (some initial results were recently presented in [4]). In our experiments, we observe that the diversity of graphs in the pre-training mixture plays an important role as well. Therefore, we believe that a brute-force increase of the model size is unlikely to bring benefits unless paired with more diverse training data and more intricate mechanisms for capturing relational interactions.

We thank the reviewer for the useful feedback and would like to comment on the weaknesses and questions.

W1. The authors have not explicitly stated the computational complexity of the method.

Short answer: The time complexity is the same as of NBFNet, i.e., it is linear in the number of edges. The memory complexity is linear in the number of nodes when using a fused message passing kernel (already provided by NBFNet) or linear in the number of edges if materializing all edge messages without the kernel.

Longer answer: The time complexity of ULTRA is upper-bounded by the entity-level GNN (because the GNN on the graph of relations has negligible overhead as the number of nodes in this graph is the same as number of unique relation types $|\mathcal{R}|$, and $|\mathcal{R}| \ll |\mathcal{V}|$ - the number of relation types is usually orders of magnitude smaller than the number of nodes, e.g, 37 relations in YAGO310 with 123k nodes). In our case, the main entity-level GNN is NBFNet, so we mainly refer to the Appendix C of the NBFNet paper [1] for all necessary derivations.

The time complexity for a single layer is generally linear in the number of edges $O(|\mathcal{E}|d + |\mathcal{V}|d^2)$. With $T$ layers, the overall complexity of a single forward pass is $O(T(|\mathcal{E}|d + |\mathcal{V}|d^2))$ but $T$ is usually a small constant (6 layers) so the complexity is essentially linear to the number of edges. However, due to the sparsity of GNNs, they are usually bounded by memory. The memory complexity of the basic NBFNet implementation is $O(T|\mathcal{E}|d)$ and linear to the number of edges, but thanks to the efficient kernelized implementation of the relational message passing (already provided by NBFNet), the memory complexity is reduced to $O(T|\mathcal{V}|d)$ and is linear in the number of nodes.

Moreover, the wall clock time and memory efficiency can be further reduced when applying more scalable and optimized versions of entity-level GNN such as AdaProp [2] or A*Net [3]. We include this discussion in the new Appendix F in the updated version of the manuscript.

Predicting all possible triples in a graph is a rather artificial task and is hardly employed even by transductive shallow embedding methods as all such methods would require computing a $O(|\mathcal{V}| \times |\mathcal{R}| \times |\mathcal{V}|)$ tensor which is rather impractical.

W2. A fair comparison would be to show the pretraining results for other inductive and transductive methods as well in addition to the SOTA comparison.

Short answer: we report the comparison with the only available pre-trainable baseline (InGram) in Table 3 dubbed as “no etypes, unconditional GNN (random)” and it is about 2x worse than ULTRA. The table caption includes the reference to InGram.

Longer answer: The term “pre-training” cannot be applied to transductive models as they cannot run inductive inference on unseen graphs because their entity/relation embedding vocabularies are fixed to that particular graph. For example, a transductive RotatE trained on FB15k237 can only run inference on FB15k237 and not on other 50+ graphs in our benchmark. We are then left with inductive models which are then categorized into two families:

• Inductive entity models, eg, NBFNet, RED-GNN, and similar) - they still learn relation embeddings pertaining to a particular data split and cannot be “pretrained” to run inference on unseen graphs, eg, NBFNet trained on FB V1 (one of the GraIL datasets from Teru et al) learns embeddings for 180 relations and can generalize to the test set of FB V1 where a graph has new entities but the same 180 relations. That said, such models still cannot generalize to unseen graphs with different relation vocabularies.
• Inductive entity and relation models, eg, InGram, ISDEA, and MTDEA are the only models known in the literature at the moment. They can generalize to unseen relation vocabularies and can be “pretrained” to run inductive inference. Among them, ISDEA and MTDEA have issues with computational complexity and cannot scale to the graphs in our benchmark (eg, doing ranking evaluation on the full entity set). The only available baseline is InGram, and we do pretrain such a model on the same pre-training mixture as ULTRA and report its performance in Table 3 under the “no etypes, unconditional GNN (random)” ablation. InGram does not use edge types in the relation graph, does not use conditional GNN encoder, and uses random Glorot initialization for relations in the relation graph. Having the same pretraining mixture, such a model significantly underperforms ULTRA averaged across 54 graphs: 0.192 MRR vs 0.366 MRR of ULTRA.

Thank you for the detailed reply and also for the additional experiments.

The responses address some of my concerns.

Before updating my score, I have further clarification questions/suggestions as below:

1. For the baselines /related works for inductive relation/link prediction there are a few additional works [1,2,3,4,5] that the authors may have missed. It would help to discuss these in the paper.

2. One concern that remains is that the method still needs to be finetuned on the entire dataset in most cases. I get that ULTRA requires 1 epoch (or thousands of steps) vs 10 epochs from scratch. But would it not be expected from a foundation model that the finetuning steps required should be orders of magnitude lesser? It would also help to have this study of the performance of ULTRA with a few examples (in addition to zero-shot performance) to understand the nature of the model in a new setting. This may also help practitioners decide how much training is required for the new dataset.

References:

1] https://proceedings.neurips.cc/paper_files/paper/2021/hash/a1c5aff9679455a233086e26b72b9a06-Abstract.html

2] https://proceedings.mlr.press/v162/yan22a.html

3] https://ieeexplore.ieee.org/abstract/document/9534355

4] https://link.springer.com/article/10.1007/s11280-023-01168-w

5] https://ojs.aaai.org/index.php/AAAI/article/view/16779

We thank the reviewer for acknowledging the merits of our work including the importance of the tackled problem, experimental results, and writing quality. Below, we would like to comment on the identified weaknesses.

If the supervised SOTA models are designed to deal with transductive settings, they may show worse performance on the downstream tasks

Yes, this is correct. To be even more precise, transductive models trained on one KG cannot run inductive inference on downstream tasks at all because their entity/relation embedding vocabularies are fixed to that particular graph. For example, a transductive RotatE trained on FB15k237 can only run inference on FB15k237 and not on other 50+ graphs in our benchmark. In that sense, the term “pre-training” cannot be applied to transductive models as they only fit one particular graph data distribution with a fixed set of entities and relations.

So, could you measure the performance of the "pretrained" SOTA models on the inductive if possible?

Short answer: we report the comparison with the only available baseline (InGram) in Table 3 dubbed as “no etypes, unconditional GNN (random)” and it is about 2x worse than ULTRA. The table caption includes the reference to InGram.

Elaborated answer: Based on the first paragraph, we cannot pretrain transductive SOTA models to run inductive inference on our benchmark. We are then left with inductive models which are then categorized into two families:

• Inductive entity models, eg, NBFNet, RED-GNN, and similar) - they still learn relation embeddings pertaining to a particular data split and cannot be “pretrained” to run inference on unseen graphs, eg, NBFNet trained on FB V1 (one of the GraIL datasets from Teru et al) learns embeddings for 180 relations and can generalize to the test set of FB V1 where a graph has new entities but the same 180 relations. That said, such models still cannot generalize to unseen graphs with different relation vocabularies.
• Inductive entity and relation models, eg, InGram, ISDEA, and MTDEA are the only models known in the literature at the moment. They can generalize to unseen relation vocabularies and can be “pretrained” to run inductive inference. Among them, ISDEA and MTDEA have issues with computational complexity and cannot scale to the graphs in our benchmark (eg, doing ranking evaluation on the full entity set). The only available baseline is InGram, and we do pretrain such a model on the same pre-training mixture as ULTRA and report its performance in Table 3 under the “no etypes, unconditional GNN (random)” ablation. InGram does not use edge types in the relation graph, does not use conditional GNN encoder, and uses random Glorot initialization for relations in the relation graph. Having the same pretraining mixture, such a model significantly underperforms ULTRA averaged across 54 graphs: 0.192 MRR vs 0.366 MRR of ULTRA.