Paper: Towards Foundation Models for Knowledge Graph Reasoning (ICLR-2024)

Reviewer comment (Novelty):
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. Such transferability has already been demonstrated by PRODIGY (https://arxiv.org/abs/2305.12600) and should be addressed.

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

Reviewer comment (Evaluation):
The model does not scale well as the authors already pointed out.

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

Reviewer comment (Writing):
Some notations are a bit hard to understand. See questions.

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

Reviewer comment (Writing):
What are u and v in h_{u|v} in section 4.2?

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

Reviewer comment (Experiments):
Why are supervised SOTA baselines only reported for some datasets in Figure 4?

Author response:
The only available baseline for WikiTopics, Metafam, and FBNELL datasets is MTDEA [3] and INDIGO [4] for BM datasets. Both models have scalability issues and only report Hits@10 against 50 randomly chosen negatives. MTDEA was evaluated for the KG completion (entity prediction) task only on some of the graphs from the benchmark, ie, Metafam, FBNELL, WikiTopics MT1:tax and MT1:health. We compare with MTDEA and INDIGO in terms of this Hits@10 (50) metric where the baseline results are available.

Since ULTRA is much more computationally efficient than both baselines, in addition to this limited metric, we also report ranking metrics (MRR and Hits@10) against the full entity set of each graph (which is a much harder task). The results indicate that evaluation against 50 random negatives overestimates the real performance and it is considerably lower when ranking against the full entity set.

Reviewer comment (Writing):
In Section 4.2, the explanation about the indicator function with variables $u$ and $v$ is a bit unclear to me. Could you elaborate more on the process and result of the indicator function according to those two variables, perhaps with visuals?

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

Reviewer comment (Experiments):
Text-based methods (e.g., LM-based methods) can be generalizable to any knowledge graphs including unseen nodes and unseen edges, as long as their nodes and edges are represented with texts. In this vein, I think one potential direction for building a foundation model for knowledge graph-related tasks 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; meanwhile, given the framing of this work ("Towards Foundation Models for Knowledge Graph Reasoning"), this point should be carefully explained.

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

Reviewer comment (Evaluation):
I would like to suggest emphasizing the performance differences between inductive and transductive setups when explaining Table 1. The proposed method w/ 0-shot settings are strong on inductive graphs; meanwhile, previous methods are superior to it on transductive graphs, which are worthwhile to discuss.

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

|  | MRR | Hits@10 |
| --- | --- | --- |
| ULTRA (8 graphs) | 0.377 | 0.537 |
| Supervised SOTA | 0.371 | 0.511 |
	
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.

Reviewer comment (Theory):
The authors have not explicitly stated the computational complexity of the method. From the paper, it seems that the forward pass is run on the entire relational graph to obtain relation representations. This is then used to initialize the node embedding from the query triple and the process is repeated for every triple. Thus it seems that the entire graph is being used for link prediction every triple making the computational complexity O(E^2). This seems limiting for large graphs that have not been explored in the paper (such as wikidata-5m etc.).

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

Reviewer comment (Experiments):
From Table 2 we can see that finetuning over the pre-trained models helps the results significantly over the 0-shot setting. Also, the fine-tuning steps are too large to claim few shot results. This weakens the claim of the "foundation model" for KGs. A fair comparison would be to show the pretraining results for other inductive and transductive methods as well in addition to the SOTA comparison.

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

Reviewer comment (Experiments):
Another limitation is that of scale. Since the current model has fewer parameters, this would limit learning over larger pretraining datasets as can be seen in Figure 6 and also reported by the authors.

Author response:
> **W3. Another limitation is that of scale. Since the current model has fewer parameters, this would limit learning over larger pretraining datasets as can be seen in Figure 6 and also reported by the authors. How would this be addressed in future works for the model?**

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 [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.

Reviewer comment (Evaluation):
SoTA results for transductive models are better than the pre-trained Ultra model in many datasets. Thus the Ultra model seems to work well for the inductive setting rather than transductive. Thus the claim of the "foundation model" seems broader in scope.

Author response:
> **W4. SoTA results for transductive models are better than the pre-trained Ultra model in many datasets. Thus the Ultra model seems to work well for the inductive setting rather than transductive. Could the authors comment on why this would be the case and how would the model be improved to handle the transductive setting?**

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 [5] and [6]. 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.

|  | MRR | Hits@10 |
| --- | --- | --- |
| ULTRA (8 graphs) | 0.377 | 0.537 |
| Supervised SOTA | 0.371 | 0.511 |
	
Of course, for those 8 graphs in the training mixture, 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 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. As we discussed in the previous paragraphs and answer to W3, the question of scaling inductive GNNs is still open and there might be many contributing factors to it.

Reviewer comment (Evaluation):
We see that in the metafam dataset, the pretraining results are poor but on finetuning the results are improved drastically. This shows that the method works well in cases where the relational patterns of the downstream datasets are similar to the pre-trained one but when the data distribution changes the results suffer. Moreover, due to limited capacity, the model may not be able to handle such cases by increasing the pretraining datasets calling for downstream finetuning. Thus domain adaptation is not a problem which can be easily overcome by scaling the current model and this further weakens the claim of a "foundation model" for KGs.

Author response:
> **W5. Metafam: the pretraining results are poor, improved after fine-tuning. The method works ... when relational patterns of the downstream datasets are similar to the pre-trained ones but otherwise suffer. why the results on this dataset are not good?**

Metafam proposed in the MTDEA paper [7] is a synthetic dataset that does not have a single connected graph in train/test splits. To be precise, the dataset has 50 disconnected components in the training split and 25 disconnected components in the test split, in both splits the size of each connected component is capped by 26 nodes. This is the only synthetic dataset among 57 graphs considered in this work whereas the other graphs are derived from more natural graphs such as Wikidata, Freebase, or DBpedia. With those reasons in mind, we would rather refrain from drawing generalization conclusions based on the single synthetic dataset in light of the results on the other 50+ graphs. As we are at the beginning of exploring transferable invariances on KGs, it is expected that no model is optimal for all tasks in the zero-shot inference mode. Still, a short fine-tuning allows ULTRA to quickly adapt to this setup with many disconnected components and solve the task almost perfectly.

Reviewer comment (Experiments):
If the supervised SOTA models are designed to deal with transductive settings, they may show worse performance on the downstream tasks.

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

Reviewer comment (Experiments):
If the SOTA models are the models for the inductive setting, they may be possible to be pretrained like the ULTRA. So, could you measure the performance of the "pretrained" SOTA models on the inductive if possible?

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