NodePiece: Compositional and Parameter-Efficient Representations of Large Knowledge Graphs (2106.12144v2)

Abstract: Conventional representation learning algorithms for knowledge graphs (KG) map each entity to a unique embedding vector. Such a shallow lookup results in a linear growth of memory consumption for storing the embedding matrix and incurs high computational costs when working with real-world KGs. Drawing parallels with subword tokenization commonly used in NLP, we explore the landscape of more parameter-efficient node embedding strategies with possibly sublinear memory requirements. To this end, we propose NodePiece, an anchor-based approach to learn a fixed-size entity vocabulary. In NodePiece, a vocabulary of subword/sub-entity units is constructed from anchor nodes in a graph with known relation types. Given such a fixed-size vocabulary, it is possible to bootstrap an encoding and embedding for any entity, including those unseen during training. Experiments show that NodePiece performs competitively in node classification, link prediction, and relation prediction tasks while retaining less than 10% of explicit nodes in a graph as anchors and often having 10x fewer parameters. To this end, we show that a NodePiece-enabled model outperforms existing shallow models on a large OGB WikiKG 2 graph having 70x fewer parameters.

• The paper proposes NodePiece, which uses an anchor-based, compositional tokenization strategy to drastically reduce parameters in large knowledge graphs.
• It employs a flexible encoder—from MLPs to Transformers—to map tokenized representations into continuous embeddings suitable for diverse downstream tasks.
• Experimental results reveal that NodePiece achieves competitive link prediction and node classification performance with up to 70 times fewer parameters.

NodePiece: Compositional and Parameter-Efficient Representations of Large Knowledge Graphs

The research presented addresses the inefficiencies within traditional knowledge graph (KG) embedding techniques, which typically involve assigning unique vectors to each entity, resulting in linear memory growth and high computational costs as the size of the graph increases. Drawing inspiration from NLP, where subword tokenization has enhanced model efficiency, this paper proposes a novel approach named NodePiece. NodePiece leverages an anchor-based mechanism that significantly reduces parameter complexity while maintaining competitive task performance.

Key Concepts and Methodology

NodePiece introduces a structured vocabulary derived from anchor nodes in a KG, allowing the model to represent entities in a compositional manner. The core idea is to replace the transductive embedding paradigm with a more flexible and parameter-efficient approach:

1. Anchor-Based Vocabulary: The method constructs a fixed-size vocabulary utilizing anchor nodes and relation types, constrained to significantly fewer parameters than the total number of graph entities.
2. Node Tokenization: Each node is represented by a combination of the $k$ nearest anchors, discrete distances to these anchors, and a relational context comprising $m$ relations. This tokenization mimics BERT's subword tokenization strategy, aiming to provide a scalable solution for new or unseen graph entities.
3. Encoding Strategy: NodePiece employs an encoder function, ranging from simple multi-layer perceptrons (MLPs) to Transformers, to map the tokenized representation into a continuous embedding space. The choice of encoder allows leveraging different inductive biases useful for varying downstream tasks.

Experimental Results

The paper evaluates NodePiece on several tasks, including transductive and inductive link prediction, as well as node classification. Results demonstrate significant parameter reductions without substantial loss of task performance:

• Transductive Link Prediction: Across various datasets, NodePiece achieves over 80-90% of the performance of state-of-the-art transductive models while utilizing up to 70 times fewer parameters—especially notable on large-scale graphs like OGB WikiKG 2.
• Inductive Learning Capabilities: For tasks requiring predictions on previously unseen entities, NodePiece shows utility without re-training or complex feature learning, performing competitively against specialized inductive models.
• Node Classification and Scalability: In node classification, NodePiece exhibits strong generalization abilities, largely reducing overfitting compared to other deep learning-based methods, thus achieving superior performance on sparse annotation scenarios.

Implications and Future Directions

This research presents theoretically grounded and empirically validated methods that shift how large-scale KGs can be efficiently represented. The implications are substantial:

• Parameter Efficiency: Reducing memory and computation requirements can lead to more sustainable deployment of models, particularly vital in resource-constrained environments or applications with fast-changing graphs.
• Graph Inductive Tasks: NodePiece's ability to generalize to inductive scenarios has significant implications for dynamic graphs found in real-world applications like social networks and recommendation systems.
• Further Exploration: Future work could focus on refining anchor selection strategies, exploring alternative encoder designs, and extending the approach to other graph types, such as hypergraphs, to capitalize on NodePiece's compositional potential.

NodePiece's framework, thus, signifies a critical stride toward more efficient, adaptable, and scalable KG embedding techniques, offering a promising avenue for ongoing research within graph representation learning.