A Foundation Model for Zero-shot Logical Query Reasoning

We thank the reviewer for appreciating our work highlighting the results, paper organization, and discussions. Please find our comments below:

W1. novelty of the proposed methods.

In this work, we introduced a new setup – a fully-inductive, zero-shot generalization of complex query answering on new, unseen KGs, and UltraQuery, the first model capable of doing that. Directly applying an Ultra model pre-trained on one-hop link prediction does not really work - the 1p query performance would be high but it is much worse on complex queries. This happens due to the multi-source propagation issue (described in Section 4.1), that is, labeling only one node is not suitable for CLQA where at intermediate stages there might be many non-zero initialized nodes representing intermediate answers. The issue can be alleviated with scores thresholding (underperforming option) or short but effective fine-tuning on a few complex queries (better performing option). Besides, we proposed 11 new inductive CLQA datasets.

Q1. I am wondering if UltraQuery can exploit the feature from entity (relation) like attributes, types, and so on to improve the performance by fine-tuning.

Yes, UltraQuery can be extended to support additional node / edge features by concatenating the features with zero/one vectors adhering to the labeling strategy in order to retain the expressive power of conditional MPNNs. However, there is a caveat when including features into the pre-trained model - since the model dimension is fixed, those features have to be homogeneous (having the same feature dimension and coming from the same distribution) while arbitrary datasets that might have textual / numerical / categorical features. So far, this feature alignment is quite non-trivial and there is no common approach for such an alignment in the literature, so we deem it an interesting avenue for the future work and will include this discussion into the relevant section.

Having said that, the hardest, most fundamental, and widely studied setting in inductive KG reasoning [1,2,3,4,5,6, and others] is generalization without features where models have to leverage the graph structure and graph inductive biases - this is what we focus on in UltraQuery.

Q2. Will you present the results of the setting provided in your further work?

We will do our best to provide initial results on some of the outlined problems in the final version. Meanwhile, we performed an experiment measuring the inductive performance as a function of several datasets in the training mixture. Please find the results in the attached PDF and in the general response.

References:
[1] Teru et al. Inductive Relation Prediction by Subgraph Reasoning. ICML 2020.
[2] Zhu et al. Neural Bellman-Ford Networks: A General Graph Neural Network Framework for Link Prediction. NeurIPS 2021.
[3] Zhang, Yao. Knowledge Graph Reasoning with Relational Digraph. WebConf 2022.
[4] Zhu et al. A*net: A scalable path-based reasoning approach for knowledge graphs. NeurIPS 2023
[5] Zhang et al. AdaProp: Learning adaptive propagation for graph neural network based knowledge graph reasoning. KDD 2023
[6] Galkin et al. Towards foundation models for knowledge graph reasoning. ICLR 2024.

We would like to thank the reviewer for the constructive feedback, please find our comments below.

W1. While the model's generalizability is a strength, the complexity introduced by the inductive reasoning might make it challenging to scale or adapt to very large KGs. The paper could have provided more details on the computational efficiency

Q1. How does ULTRAQUERY handle extremely large KGs, and what are the scalability challenges?

Given that there haven’t been any inductive models for CLQA, we first focused on the general problem of making zero-shot inductive transfer possible. Scalability, although being somewhat orthogonal to the inductive inference and generalization problems, is a nice bonus for inductive models once inductiveness is achieved and we believe it is an encouraging direction for future work.

Having said that, compared to the transductive baseline GNN-QE, the complexity overhead induced by the fully-inductive relation projection operator is rather negligible. The fully-inductive relation projection operator runs over the graph of relations and is O(|R|) where |R| is the number of unique relations. Usually |R| is orders of magnitude smaller than the number of nodes, eg, |R|=474 in FB15k237 with 15k nodes, or |R|=400 in NELL995 with 63k nodes.

The support for very large KGs could be further improved by adopting more scalable entity-level GNN predictors like A*Net [1] or AdaProp [2] which have been shown to scale to graphs of millions of nodes. We are optimistic that UltraQuery could scale to such graphs as well when integrated with those models. We deem it as the next engineering step for future work and will include this discussion into the revised version.

W2. The paper does not present formal theoretical results or proofs to support the empirical findings, which could have strengthened the contribution

We do not see why it is necessarily a weakness. Our model coincides with the most recent theoretical results in the relational GNN expressivity literature, eg, relational WL [3] and its extension to conditional MPNNs [4] paved the way for generalizable inductive models which are provably more powerful than shallow transductive embedding approaches. The product fuzzy logic we use is identified as the most stable fuzzy logic for backprop in the differentiable fuzzy logic literature [5]. Finally, theoretical expressivity of the inductive projection operator (Ultra) is still an open question and we are not aware of the formal theoretical results about it in the literature. We believe this an intriguing question for the future work.

W3. The reliance on a pre-trained model could lead to overfitting on the training dataset, potentially affecting the model's performance on unseen data.

Could you please elaborate on your understanding of overfitting in this case? We do not observe strong signs of overfitting on the training dataset as UltraQuery generalizes in the zero-shot manner to 20+ unseen graphs (and queries over those graphs) often better than tailored transductive models trained specifically on each target graph. In that sense, transductive models are extreme cases of overfitting as they are hardcoded to a particular set of entities or relations in a specific KG and cannot generalize to unseen graphs at inference time.

What we do observe is the multi-source propagation issue (explained in Section 4.1), which might be understood as a kind of overfitting on 1p queries when using pre-trained KG completion model for CLQA. We have discussed this issue and proposed two solutions: scores thresholding and short fine-tuning.

Generally, we would expect that a better pre-trained relation projection operator (with better zero-shot generalization capabilities on one-hop link prediction) would likely yield a higher performance in the UltraQuery framework as well (after doctoring the multi-source propagation issue).

Q2. Can the authors provide more insights into the decision-making process behind the choice of fuzzy logics for logical operations?

Good differentiable fuzzy logics must be stable in the gradient-based optimization setup (like backpropagation) and do not vanish the gradient. Van Krieken et al [5] have shown that only product logic satisfies those requirements whereas other options like Gödel and Lukasiewicz t-norms and t-conorms suffer from the vanishing gradient problem. We included the main motivation in Section 4.2 and will elaborate on that in the revised version.

References:
[1] Zhu et al. A*net: A scalable path-based reasoning approach for knowledge graphs. NeurIPS 2023
[2] Zhang et al. AdaProp: Learning adaptive propagation for graph neural network based knowledge graph reasoning. KDD 2023
[3] Barcelo et al. Weisfeiler and Leman go relational. LoG 2022
[4] Huang et al. A theory of link prediction via relational Weisfeiler-Leman on knowledge graphs. NeurIPS 2023
[5] van Krieken et al. Analyzing differentiable fuzzy logic operators. Artificial Intelligence, 302, 2022

We thank the reviewer for appreciating our work and would like to comment on the weaknesses:

W1. I feel the main problem is the lack of a detailed description of the ULTRA method and how it is adapted to the foundation model.

In order to build a model that zero-shot generalizes to CLQA tasks on any unseen KG we need two main components that do not depend on input entity/relation vocabulary: (1) generalizable logical operators; (2) generalizable relation projection operator. (1) is achieved by using non-parametric fuzzy logics (product logic in our case as being the most stable for gradient-based optimization). (2) is harder and requires an inductive link predictor / KG completion operator that generalizes to any graph. We integrated ULTRA as the projection operator that outputs a single scalar for each node (which is then used by fuzzy operators) that indicates a probability of this node being the answer to the query $(h, q, ?)$.

Elaborated in Section 4.1, practically, ULTRA consists of two GNNs – operating on the relation level and on the entity level. First, given a query $(h, q, ?)$, ULTRA builds a graph of relations from the original KG - this graph is small and only has O(|R|) nodes and 4 specific meta-relations. Note that $|R| << |E|$ in any KG so processing this graph introduces a rather marginal computational overhead compared to the entity-level GNN. Since we know the query relation $q$, we label it with the all-ones vector in the graph of relations, run a GNN, and read all final representations as conditionally dependent on the starting query $q$. Those representations are used as initial edge type features in the second, entity-level GNN. Labeling the starting head node $h$ with the vector corresponding to $q$ and running a GNN over the main graph, we read out final node states as probabilities of each node being the answer to the query $(h, q, ?)$.

The only learnable parameters in ULTRA are 4 meta-relations for the graph of relations and GNN weights. The 4 meta-relations represent structural patterns like “a tail of and edge with relation X is the head of another edge with relation Y” and can be mined from any multi-relational KG independent of their entity/relation vocabulary. GNN weights are optimized during pre-training. Since the model does not rely on any KG-specific entity or relation vocabulary, a single pre-trained ULTRA model can be used for zero-shot inference on any KG.

Applying ULTRA for CLQA in this work, we found the multi-source propagation issue (Section 4.1), that is, ULTRA pre-trained with a single starting node, is not suitable for CLQA where at intermediate stages there might be many non-zero initialized nodes representing intermediate answers. The issue can be alleviated with scores thresholding (underperforming option) or short but effective fine-tuning on a few complex queries (better performing option). Finally, we proposed 11 new inductive CLQA datasets.

Thanks for highlighting the lack of a description – we will include this more detailed discussion into the main body of the manuscript in the final version which allows one more extra page of content.

Thank you for appreciating our work and helpful comments.

W1. A more in-depth analysis of the model's behavior under different conditions (e.g., varying graph sizes, query complexity) could provide further insights

We would like to point your attention to Section 5.3 in the main paper and Appendix C. Appendix C provides breakdown of the zero-shot performance w.r.t. different query types and different graph sizes. Section 5.3 includes a qualitative analysis on faithfulness as a function of varying graph sizes, and cardinality prediction.

In Section 5.3 we study faithfulness (the ability to recover easy answers achievable by graph traversal, visualized as a function of increasing unseen inference graph size from inductive FB datasets) and cardinality prediction (correlation between model outputs and the true number of query answers) and show that UltraQuery performs quite competitively compared to larger transductive SOTA models while being inductive and orders of magnitude smaller.

In Appendix C we plot per-query-type performance on all 14 query types across 9 increasingly larger inference graphs - showing that UltraQuery performance does not degrade on more complex patterns (such as 3p or inp) when an unseen inference graph becomes larger. In contrast, the performance of baseline GNN-QE degrades quite significantly.

W2. The paper does not discuss the model's scalability to very large KGs, which is an important aspect for practical applications.

Thanks for bringing this question up. We will include this discussion into the future work. In this work, we first focused on the general ability to perform CLQA on any KG and scalability is the natural next step. To this end, any scalable path-based GNN like A*Net [1] or AdaProp [2] could be used as an entity-level GNN in the relation projection operator thus making it more of an engineering task. Since those GNNs can scale to graphs with millions of nodes, we are optimistic that UltraQuery could scale to such sizes as well.

W3. Scalability of the model (parameters and performance) can be further discussed.

When training from scratch, we found marginal performance improvements beyond a hidden dimension of 64. When using a pre-trained relation projection checkpoint, we attribute most of the performance to the baseline performance of the checkpoint - since available checkpoints of Ultra are of the same parameter size, we only used the available one. We note that most models on knowledge graphs are designed for a single size, including the checkpoint of Ultra we use. There is no consensus in the community on what should be scaled.

Q1. How the model scales as the pretraining data scales up?

We ran new experiments training UltraQuery on 2 (FB15k237, NELL995) and 3 (FB15k, NELL995, FB15k) CLQA datasets and measuring inductive inference performance on the rest 20 inductive datasets. Please find the charts and discussion in the attached PDF document and in the general response.

References:
[1] Zhu et al. A*net: A scalable path-based reasoning approach for knowledge graphs. NeurIPS 2023
[2] Zhang et al. AdaProp: Learning adaptive propagation for graph neural network based knowledge graph reasoning. KDD 2023