PaperlayerView exact revision on arXiv

Source-linked AI summary

Towards Foundation Models for Knowledge Graph Reasoning

Mikhail Galkin, Xinyu Yuan, Hesham Mostafa, Jian Tang, Zhaocheng Zhu

arXiv:2310.04562v2cs.CLcs.AI

TL;DR

Knowledge graphs need transferable representations because their entity and relation vocabularies differ across graphs. ULTRA conditions relational representations on learned interactions, enabling zero-shot transfer and fine-tuning across unseen KGs. Across 57 KGs, a single pre-trained model often matches or exceeds strong graph-specific baselines, with fine-tuning improving performance further.

  • Problem

    Knowledge graphs lack shared entity and relation vocabularies, limiting foundation-model transfer to arbitrary graphs.

    Source: §1 · p. 2

  • Method

    ULTRA builds a graph of relation interactions and uses a labeled GNN to produce relative representations without graph-specific embeddings or input features.

    Source: §1 · p. 2, §4.1 · p. 5

  • Results

    A single pre-trained ULTRA model outperforms tailored supervised baselines on 50+ graphs by average 15% in zero-shot inference, while fine-tuning adds 10% average improvement.

    Source: §6 · p. 9

  • Takeaways & Limitations

    ULTRA supports zero-shot and fine-tuned KG reasoning across arbitrary graph sizes and relation vocabularies using transferable relational structure.

    Source: §1 · p. 2, §6 · p. 9

    More pre-training graphs do not consistently improve inference, possibly because the 177k-parameter model has limited capacity.

    Source: §6 · p. 9

Abstract

from arXiv · show

Foundation models in language and vision have the ability to run inference on any textual and visual inputs thanks to the transferable representations such as a vocabulary of tokens in language. Knowledge graphs (KGs) have different entity and relation vocabularies that generally do not overlap. The key challenge of designing foundation models on KGs is to learn such transferable representations that enable inference on any graph with arbitrary entity and relation vocabularies. In this work, we make a step towards such foundation models and present ULTRA, an approach for learning universal and transferable graph representations. ULTRA builds relational representations as a function conditioned on their interactions. Such a conditioning strategy allows a pre-trained ULTRA model to inductively generalize to any unseen KG with any relation vocabulary and to be fine-tuned on any graph. Conducting link prediction experiments on 57 different KGs, we find that the zero-shot inductive inference performance of a single pre-trained ULTRA model on unseen graphs of various sizes is often on par or better than strong baselines trained on specific graphs. Fine-tuning further boosts the performance.

1 INTRODUCTION

Knowledge graphs lack the shared entity and relation vocabularies that support foundation-model transfer, motivating ULTRA’s relationally conditioned representations for arbitrary unseen graphs.

  • KGs have distinct entity and relation vocabularies, limiting transfer beyond graph-specific embeddings and fixed-relation inductive methods.Source: §1 · p. 2
  • ULTRA learns transferable invariances from interactions among relations rather than graph-specific vocabularies or input features.It constructs a relation graph and applies a labeled GNN to obtain relative representations for unseen relations.Source: §1 · p. 2
  • 50+ unseen KGs spanning 1,000–120,000 nodes and 5K–1M edges were evaluated using ULTRA pre-trained on three KGs.Source: §1 · p. 2
  • Zero-shot performance might exceed strong supervised baselines by up to 300%, while short fine-tuning often improves performance further.Source: §1 · p. 2

2 RELATED WORK

Prior inductive methods generally handle unseen entities but not unseen relations, while approaches supporting unseen relations face scalability, featurization, or language-transfer constraints.

  • Most inductive link-prediction methods generalize to new nodes only because their node representations depend on a fixed relation vocabulary.Source: §2 · p. 2
  • Meta-learning and few-shot methods support unseen relations but are computationally expensive, difficult to scale, or rely on sampled support sets.Source: §2 · p. 2
  • RMPI and InGram use relation graphs, but RMPI has scalability issues and InGram’s degree-based features transfer only to similarly distributed graphs.Source: §2 · pp. 2–3
  • ISDEA and MTDEA provide equivariant transfer to unseen relations but are too computationally expensive for the graphs studied here.Source: §2 · p. 3
  • ULTRA uses structural information without input features and is presented as scalable to millions of edges with marginal inference overhead.Source: §2 · p. 3

3 PRELIMINARIES

Fully-inductive KG reasoning requires generalizing to new entities and relation types, while link-prediction GNNs use conditional representations and ULTRA’s relation graph captures transferable interactions.

  • The hardest inductive setting has disjoint training and inference entity and relation sets, and includes the easier fixed-relation setting.Source: §3 · p. 3
  • KG reasoning predicts missing targets in queries such as (h, r, ?) or (?, r, t), with inverse relations rewriting head queries.Source: §3 · p. 3
  • NBFNet conditions node representations on the query relation, helping break neighborhood symmetries that standard GNN features cannot distinguish.Source: §3 · p. 3
  • ULTRA’s relation graph represents interactions among relation types so relative relation representations can generalize to new relations and entities.Source: Figure 2 · p. 4

4 METHOD

ULTRA enables inductive reasoning across graphs with unseen entities and relations by representing relation interactions rather than relying on fixed vocabularies. Its method constructs a relation graph, derives query-conditioned relation representations, and uses them in entity-level link prediction.

4 METHOD

  • ULTRA represents unseen relations through invariant interactions in a relation graph, enabling reasoning when training and inference relation vocabularies differ.The four interactions are tail-to-head, head-to-head, head-to-tail, and tail-to-tail; relation representations are conditioned on the query relation.Source: §4 · p. 4

4.1 RELATION GRAPH CONSTRUCTION

  • The relation graph contains one node per relation type and four typed edge interactions, forming the structural basis for transferable relational representations.The full relation-graph adjacency has four channels, one for each fundamental interaction.Source: §4.1 · p. 5

4.2 CONDITIONAL RELATION REPRESENTATIONS

  • ULTRA initializes the queried relation in the relation graph and applies message passing to obtain query-conditioned representations for every original relation.The conditional representations have dimension d and are produced by a relation-level GNN using the four interaction embeddings and update layers.Source: §4.2 · p. 5

4.3 ENTITY-LEVEL LINK PREDICTION

  • The entity-level predictor initializes the query head with its relation representation, propagates messages using transformed relation features, and scores candidate tail nodes.A final MLP maps node states to logits for the query, while the predictor can use off-the-shelf inductive link-prediction architectures that require relational features.Source: §4.3 · pp. 5–6
  • ULTRA is trained with binary cross-entropy over positive triples and negatives generated by corrupting either the head or tail.The framework can be trained on any multi-relational graph or mixture of graphs.Source: §4.3 · p. 6

5 EXPERIMENTS

ULTRA is evaluated across 57 knowledge graphs for zero-shot transfer, fine-tuning, and ablations of its transferable relation representations. Across the benchmark, zero-shot performance is competitive with supervised baselines, fine-tuning improves results, and conditional relation encoding supports generalization.

5.1 SETUP AND DATASETS

  • ULTRA is evaluated on 57 knowledge graphs spanning transductive, inductive-entity, and inductive-entity-and-relation settings, with graph sizes ranging from 1k to 120k entities.The evaluation includes zero-shot inference, fine-tuning, and comparisons with dataset-specific baselines.Source: §5.1 · p. 6, §5 · p. 6

5.2 MAIN RESULTS: ZERO-SHOT INFERENCE AND FINE-TUNING OF ULTRA

  • Zero-shot ULTRA outperforms supervised baselines on average in both MRR and Hits@10, with gains approaching 3× on smaller inductive graphs such as FB-25 and FB-50.On FB-25 and FB-50, the reported relative gains are 291% and 289%, respectively.Source: §5.2 · p. 7, Table 1 · p. 7, Table 2 · p. 8
  • Fine-tuning further improves ULTRA, increasing average MRR by 10% across 54 graphs and especially helping on larger transductive datasets.Fine-tuning ULTRA typically requires 1000–2000 steps, although some datasets show marginal or negative changes.Source: §5.2 · p. 8, Table 2 · p. 8

5.3 ABLATION STUDY

  • A single pre-trained ULTRA model performs almost on par with models trained separately on each graph, while fine-tuning achieves the overall best performance with fewer computational resources.Zero-shot ULTRA lags on larger transductive graphs but performs better on inductive datasets.Source: §5.3 · p. 8, Figure 5 · p. 8
  • Adding more than three graphs to the pre-training mixture produces saturated inductive inference performance, with model capacity, scale, and optimization left for future study.The main model is pre-trained on three knowledge graphs.Source: §5.1 · p. 6, §5.3 · pp. 8–9
  • Removing relation edge types and using an unconditional relation encoder causes up to a 48% relative MRR drop, supporting conditional relation and entity representations for transferable link prediction.The cited comparison is 0.192 versus 0.366 MRR for the homogeneous relation graph with randomly initialized node states.Source: §5.3 · p. 9, Table 3 · p. 9

6 DISCUSSION AND FUTURE WORK

ULTRA is presented as a transferable approach for KG reasoning, with reported zero-shot and fine-tuned gains across diverse graphs. The section also identifies scaling limitations, broader efficiency and misuse implications, and reproducibility resources.

6 DISCUSSION AND FUTURE WORK

  • More pre-training graphs do not often improve inference, potentially because the 177k-parameter model has limited capacity; scaling beyond 200k parameters also showed no significant preliminary gains.The authors hypothesize input normalization and optimization may contribute to the unresolved scaling behavior.Source: §6 · p. 9
  • A single pre-trained ULTRA model outperforms tailored supervised baselines on 50+ graphs by average 15% in zero-shot evaluation, while fine-tuning adds a further 10% average improvement.The conclusion covers graphs with 1k–120k nodes.Source: §6 · p. 9
  • ULTRA supports training and inference on multi-relational graphs without input features by using relational-structure invariance and conditional relation representations.Source: §6 · p. 9

ETHICS STATEMENT

  • Foundation graph models could reduce the computational burden and carbon footprint of training graph-specific models, while also enabling misuse such as searching anonymized data for patterns.Source: ETHICS STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

  • The paper provides dataset and evaluation details, hyperparameters, full MRR and Hits@10 results with standard deviations, and source code in supplementary materials.Source: REPRODUCIBILITY STATEMENT · p. 10

A DATASETS

The evaluation spans 57 openly available KGs across transductive and two inductive settings. Dataset-specific published state-of-the-art models are aggregated as supervised baselines, while saturated small datasets are excluded.

A DATASETS

  • Evaluation covers 57 KGs divided into 16 transductive, 18 inductive-entity, and 23 inductive entity-and-relation datasets.Source: A DATASETS · p. 16
  • Each dataset is compared with a published model trained specifically on that target graph, whose performance is aggregated as Supervised SOTA.Source: A DATASETS · p. 16
  • Kinships, UMLS, Countries, and Family are omitted because their performance is saturated and considered non-representative.Source: A DATASETS · p. 16
  • For HM 1k, HM 3k, and HM 5k, the evaluation predicts both heads and tails and compares against the corresponding reported baselines.Source: A DATASETS · p. 16
  • MT2, MT3, and MT4 lack reported entity-only completion results, while Hetionet uses TorchDrug splits and its reported RotatE baseline.Source: A DATASETS · p. 16

B SPARSE MMS FOR RELATION GRAPH

The relation graph is constructed from the original multi-relational graph using sparse matrix multiplications over head-relation and tail-relation incidence matrices. Stacking four interaction adjacencies yields the relation graph with fundamental edge types.

B SPARSE MMS FOR RELATION GRAPH

  • Four sparse matrix multiplications represent the fundamental relation types h2t, h2h, t2h, and t2t.Source: B SPARSE MMS FOR RELATION GRAPH · p. 16
  • The method builds sparse head-relation and tail-relation matrices from the original adjacency tensor, then computes relation interactions through sparse matrix multiplication.Eh captures head-relation pairs and Et captures tail-relation pairs.Source: B SPARSE MMS FOR RELATION GRAPH · p. 16
  • The nonzero entries of the four sparse matrices define edge indices that are stacked into the final relation-graph adjacency tensor and graph.Source: B SPARSE MMS FOR RELATION GRAPH · p. 16

C HYPERPARAMETERS

The appendix specifies ULTRA’s pre-training configuration, dataset table scopes, and procedures for fine-tuning, repeated runs, and graph-mixture ablations.

C HYPERPARAMETERS

  • ULTRA pre-training uses 6-layer GNNs with hidden dimension 64, DistMult messages, sum aggregation, and a 2-layer MLP for conditional relation transformations.Source: C HYPERPARAMETERS · p. 16
  • The model is trained for 200,000 steps on FB15k237, WN18RR, and CoDEx-Medium using AdamW, learning rate 0.0005, and batch size 64.Source: C HYPERPARAMETERS · p. 16
  • Tables 4–6 document transductive, inductive-entity, and inductive entity-and-relation datasets, while Table 7 lists GNNr and GNNe pre-training hyperparameters.Source: Table 4 · p. 17, Table 5 · p. 17, Table 6 · p. 17
  • Fine-tuning selects the best validation checkpoint when using multiple epochs, with five random-seed repetitions; models trained from scratch use three repetitions.Source: C HYPERPARAMETERS · pp. 17–18
  • Graph-mixture ablations reduce training steps for mixtures of five or more large graphs to finish within three days on two A100 GPUs.Source: C HYPERPARAMETERS · p. 18

D FULL RESULTS

Across the full evaluation, ULTRA is compared with supervised baselines across diverse graph settings, while results vary with graph scale, sparsity, relation count, and pre-training diversity.

D FULL RESULTS

  • Full results cover 43 graphs with published SOTA comparisons and 16 additional graphs with partially reported Hits@10 estimates.Table 10 spans pre-training, inductive entity, inductive entity-and-relation, and transductive graphs; Table 11 uses 50 randomly chosen negatives where available.Source: D FULL RESULTS · p. 18
  • The 50-negative Hits@10 metric can greatly overestimate performance, motivating full MRR and Hits@k evaluation against the whole entity set.The authors recommend reporting whole-entity-set metrics in future work.Source: D FULL RESULTS · p. 18
  • The largest performance gaps occur on larger, highly relational, or very sparse graphs outside the medium-sized three-graph training mixture.The cited examples include graphs with up to 120k nodes, 1M edges, more than 1,600 relation types, or extreme sparsity.Source: D FULL RESULTS · p. 18
  • 0.377 MRR / 0.537 Hits@10 versus 0.371 MRR / 0.511 Hits@10: an ULTRA checkpoint pre-trained on 8 graphs outperforms supervised SOTA on 16 transductive graphs.This comparison uses the expanded pre-training mixture rather than the three-graph checkpoint.Source: D FULL RESULTS · p. 19
  • ULTRA checkpoints contain about 170k parameters, whereas transductive baselines can use many millions, including 80M parameters for BiQUE on ConceptNet100k.Transductive models can memorize target-specific distributions, but that parameterization does not transfer across KGs.Source: D FULL RESULTS · p. 19
  • The authors report no established GNN scaling laws and argue that increasing model size alone may not help without more diverse training data and better relational-interaction mechanisms.They also identify pre-training mixture diversity as important in their experiments.Source: Table 9 · p. 19, D FULL RESULTS · p. 20

E ON ADDING MORE FEATURES

ULTRA focuses on structural representations that apply to featureless or feature-augmented KGs, while dataset-specific text or numerical features may provide complementary gains.

E ON ADDING MORE FEATURES

  • ULTRA focuses on structural representations and feature-less graphs so it can apply to KGs with or without heterogeneous domain-specific features.Different domains may provide non-overlapping biomedical, geographical, numerical, or textual attributes requiring distinct encoders.Source: E ON ADDING MORE FEATURES · p. 20
  • Concatenating encoded text features with structural GNN features may further improve inductive performance, especially during fine-tuning.The authors treat dataset-specific features as complementary to ULTRA representations and leave this combination for future work.Source: E ON ADDING MORE FEATURES · p. 20

F COMPUTATIONAL COMPLEXITY

ULTRA’s computational cost is governed mainly by the entity-level GNN, with near-linear dependence on graph size and memory that can scale linearly with nodes.

F COMPUTATIONAL COMPLEXITY

  • ULTRA’s time complexity is dominated by the entity-level GNNe because the relation-graph GNN has negligible overhead when |R| ≪ |V|.The main entity-level GNNe is NBFNet in the reported implementation.Source: F COMPUTATIONAL COMPLEXITY · p. 20
  • O(T(|E|d+|V|d2)): a T-layer forward pass is essentially linear in edges when T is a small constant such as 6.The single-layer cost is O(|E|d + |V|d2).Source: F COMPUTATIONAL COMPLEXITY · p. 20
  • Efficient kernelized message passing reduces NBFNet memory complexity from O(T|E|d) to O(T|V|d), making memory linear in the number of nodes.More scalable entity-level GNNs such as AdaProp or A*Net can reduce complexity further.Source: F COMPUTATIONAL COMPLEXITY · p. 20
  • Tables 10 and 11 report ULTRA’s full MRR and Hits@10 results in zero-shot and fine-tuning regimes across the evaluated graph sets.These tables provide the empirical evaluation associated with the computationally described model.Source: Table 10 · p. 21, Table 11 · p. 22
Loading 2310.04562v2…