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

1 INTRODUCTION . Representation learning tasks on knowledge graphs ( KGs ) often require a parameterization of each unique atom in the graph with a vector or matrix . Traditionally , in multi-relational KGs such atoms constitute a set of all nodes n ∈ N ( entities ) and relations ( edge types ) r ∈ R ( Nickel et al. , 2016 ) . Assuming parameterization with vectors , atoms are mapped to d-dimensional vectors through shallow encoders fn : n → Rd and fr : r → Rd which scale linearly to the number of nodes and edge types2 , i.e. , having O ( |N | ) space complexity of the entity embedding matrix . Albeit efficient on small conventional benchmarking datasets based on Freebase ( Toutanova & Chen , 2015 ) ( ~15K nodes ) and WordNet ( Dettmers et al. , 2018 ) ( ~40K nodes ) , training on larger graphs ( e.g. , YAGO 3-10 ( Mahdisoltani et al. , 2015 ) of 120K nodes ) becomes computationally challenging . Scaling it further up to larger subsets ( Hu et al. , 2020 ; Wang et al. , 2021 ; Safavi & Koutra , 2020 ) of Wikidata ( Vrandecic & Krötzsch , 2014 ) requires a top-level GPU or a CPU cluster as done in , e.g. , PyTorch-BigGraph ( Lerer et al. , 2019 ) that maintains a 78M × 200d embeddings matrix in memory ( we list sizes of current best performing models in Table 1 ) . Taking the perspective from NLP , shallow node encoding in KGs corresponds to shallow word embedding popularized with word2vec ( Mikolov et al. , 2013 ) and GloVe ( Pennington et al. , 2014 ) that learned a vocabulary of 400K-2M most frequent words , treating rarer ones as out-of-vocabulary ( OOV ) . The OOV issue was resolved with the ability to build infinite combinations with a finite vocabulary enabled by subword units . Subword-powered algorithms such as fastText ( Bojanowski et al. , 2017 ) , Byte-Pair Encoding ( Sennrich et al. , 2016 ) , and WordPiece ( Schuster & Nakajima , 2012 ) became a standard step in preprocessing pipelines of large language models and allowed to construct fixed-size token vocabularies , e.g. , BERT ( Devlin et al. , 2019 ) contains ~30K tokens and 1The code is available on GitHub : https : //github.com/migalkin/NodePiece 2We then concentrate on nodes as usually their size is orders of magnitude larger than that of edge types . GPT-2 ( Radford et al. , 2019 ) employs ~50K tokens . Importantly , relatively small input embedding matrices enabled investing the parameters budget into more efficient encoders ( Kaplan et al. , 2020 ) . Drawing inspiration from subword embeddings in NLP , we explore how similar strategies for tokenizing entities in large graphs can dramatically reduce parameter complexity , increase generalization , and naturally represent new unseen entities as using the same fixed vocabulary . To do so , tokenization has to rely on atoms akin to subword units and not the total set of nodes . To this end , we propose NodePiece , an anchor-based approach to learn a fixed-size vocabulary V ( |V | |N | ) of any connected multi-relational graph . In NodePiece , the set of atoms consists of anchors and all relation types that , together , allow to construct a combinatorial number of sequences from a limited atoms vocabulary . In contrast to shallow approaches , each node n is first tokenized into a unique hash ( n ) of k closest anchors and m immediate relations . A key element to build a node embedding is a proper encoder function enc ( n ) : hash ( n ) → Rd which can be designed leveraging inductive biases of an underlying graph or downstream tasks . Therefore , the overall parameter budget is now defined by a small fixed-size vocabulary of atoms and the complexity of the encoder function . Our experimental findings suggest that a fixed-size NodePiece vocabulary paired with a simple encoder still yields competitive results on a variety of tasks including link prediction , node classification , and relation prediction . Furthermore , anchor-based hashing enables conventional embedding models to work in the inductive and out-of-sample scenarios when unseen entities arrive at inference time , which otherwise required tailored learning mechanisms . 2 RELATED WORK . Conventional KG embedding approaches . To the best of our knowledge , all contemporary embedding algorithms ( Ji et al. , 2020 ; Ali et al. , 2020 ) for link prediction on KGs employ shallow embedding lookups mapping each entity to a unique embedding vector thus being linear O ( |N | ) to the total number of nodes |N | and size of an embedding matrix . This holds for different embedding families , e.g. , translational ( Sun et al. , 2019 ) , tensor factorization ( Lacroix et al. , 2018 ) , convolutional ( Dettmers et al. , 2018 ) , and hyperbolic ( Chami et al. , 2020 ; Balazevic et al. , 2019 ) . The same applies to relation-aware graph neural network ( GNN ) encoders ( Schlichtkrull et al. , 2018 ; Vashishth et al. , 2020 ) who still initialize each node with a learned embedding or feature vector before message passing . Furthermore , shallow encoding is also used in higher-order KG structures such as hypergraphs ( Fatemi et al. , 2020 ) and hyper-relational graphs ( Rosso et al. , 2020 ; Galkin et al. , 2020 ) . NodePiece can be used as a drop-in replacement of the embedding lookup with any of those models . Distillation and compression . Several recent techniques for reducing memory footprint of embedding matrices follow successful applications of distilling large language models in NLP ( Sanh et al. , 2019 ) , i.e. , distillation ( Wang et al. , 2020 ; Zhu et al. , 2020 ) into low-dimensional counterparts , and compression of trained matrices into discrete codes ( Sachan , 2020 ) . However , all of them require a full embedding matrix as input which we aim to avoid designing NodePiece . Vocabulary reduction in recommender systems . Commonly , recommender systems operate on thousands of categorical features combined in sparse high-dimensional vectors.Recent approaches ( Medini et al. , 2021 ; Liang et al. , 2021 ) employ anchor-based hashing techniques to factorize sparse feature vectors into dense embeddings . Contrary to those setups , we do not expect availability of feature vectors for arbitrary KGs and rather learn vocabulary embeddings from scratch . Entity descriptions and language models . A recent line of work such as KG-BERT ( Yao et al. , 2019 ) , MLMLM ( Clouâtre et al. , 2021 ) , BLP ( Daza et al. , 2021 ) utilize entity descriptions passed through a language model ( LM ) encoder as entity embeddings suitable for link prediction . We would like to emphasize that such approaches are rather orthogonal to NodePiece . Textual features are mostly available in Wikipedia-derived KGs like Wikidata but are often missing in domain-specific graphs like social networks and product graphs . We therefore assume textual features are not available and rather learn node representations based on their spatial characteristics . Still , textual features can be easily added by concatenating NodePiece-encoded features with LM-produced features . Out-of-sample representation learning . This task focuses on predictions involving previously unseen , or out-of-sample , entities that attach to a known KG with a few edges . These new edges are then utilized as a context to compute its embedding . Previous work ( Wang et al. , 2019 ; Hamaguchi et al. , 2017 ; Albooyeh et al. , 2020 ) proposed different neighborhood aggregation functions for this process or resorted to meta-learning ( Chen et al. , 2019 ; Baek et al. , 2020 ; Zhang et al. , 2020a ) . However , all of them follow the shallow embedding paradigm . Instead , NodePiece uses the new edges as a basis for anchor-based tokenization of new nodes in terms of an existing vocabulary . 3 NODEPIECE VOCABULARY CONSTRUCTION . Given a directed KG G = ( N , E , R ) consisting of |N | nodes , |E| edges , and |R| relation types , our task is to reduce the original vocabulary size of |N | nodes to a smaller , fixed-size vocabulary of node pieces akin to subword units . In this work , we represent node pieces through anchor nodes a ∈ A , A ⊂ N , a pre-selected set of nodes in a graph following a deterministic or stochastic strategy . A full NodePiece vocabulary is then constructed from anchor nodes and relation types , i.e , V = A+R . Note that in order to maintain reachability of each node and balance in- and out-degrees we enrich G with inverse edges with inverse relation types , such that |R|inverse = |R|direct and |R| = |R|direct + |R|inverse . Using elements of the constructed vocabulary each node n can be tokenized into hash ( n ) as a sequence of k closest anchors , discrete anchor distances , and a relational context of m immediate relations . Then , any encoder function enc ( n ) : hash ( n ) → Rd can be applied to embed the hash into a d-dimensional vector . An intuition of the approach is presented in Fig . 1 with each step explained in more detail below . 3.1 ANCHOR SELECTION . Subword tokenization algorithms such as BPE ( Sennrich et al. , 2016 ) employ deterministic strategies to create tokens and construct a vocabulary , e.g. , based on frequencies of co-occurring n-grams , such that more frequent words are tokenized with fewer subword units . On graphs , such strategies might employ centrality measures like degree centrality or Personalized PageRank ( Page et al. , 1999 ) . However , in our preliminary experiments , we found random anchor selection to be as effective as centrality-based strategies . A choice for deterministic strategies might be justified when optimizing for certain task-specific topological characteristics , e.g. , degree and PPR strategies indeed skew the distribution of shortest anchor distances towards smaller values thus increasing chances to find anchors in 2- or 3-hop neighborhood of any node ( we provide more evidence for that in Appendix C ) . 3.2 NODE TOKENIZATION . Once the vocabulary V = A + R is constructed , each node n can be hashed ( or tokenized ) into a hash ( n ) using 1 ) k nearest anchors and their discrete distances ; 2 ) m immediate outgoing relations from the relational context of n. Since anchor nodes are concrete nodes in G , they get hashed in the same way as other non-anchor nodes . Anchors per node . Given |A| anchor nodes , it is impractical to use all of them for encoding each node . Instead , we select k anchors per node and describe two possible strategies for that , i.e. , random and deterministic . The basic random strategy uniformly samples an unordered set of k anchors from A yielding ( |A| k ) possible combinations . To avoid collisions when hashing the nodes , |A| and k are to be chosen according to the lower bound on possible combinations that is defined by the total number of nodes , e.g. , ( |A| k ) ≥ |N | . Note that running depth-first search ( DFS ) to random anchors at inference time is inefficient and , therefore , hash ( n ) of the random strategy has to be pre-computed . On the other hand , the deterministic strategy selects an ordered sequence of k nearest anchors . Hence , the anchors can be obtained via breadth-first search ( BFS ) in the l-hop neighborhood of n at inference time ( or pre-computed for speed reasons ) . However , the combinatorial bound is not applicable in this strategy and we need more discriminative signals to avoid hash collisions since nearby nodes will have similar anchors ( we elaborate on the uniqueness issue in Appendix K ) . Such signals have to better ground anchors to the underlying graph structure , and we accomplish that using anchor distances3 and relational context described below . A node residing in a disconnected component is assigned with an auxiliary [ DISCONNECTED ] token or can be turned into an anchor . However , the majority of existing KGs are graphs with one large connected component with very few disconnected nodes , such that this effect is negligible . Anchor Distances . Given a target node n and an anchor ai , we define anchor distance zai ∈ [ 0 ; diameter ( G ) ] as an integer denoting the shortest path distance between ai and n in the original graph G. Note that when tokenizing an anchor aj with the deterministic strategy , the nearest anchor among top-k is always aj itself with distance 0 . We then map each integer to a learnable d-dimensional vector fz : zai → Rd akin to relative distance encoding scheme . Relational Context . We also leverage the multi-relational nature of an underlying KG . Commonly4 , the amount of unique edge types in G is orders of magnitude smaller than the total number of nodes , i.e. , |R| |N | . This fact allows to include the entire |R| in the NodePiece vocabulary VNP and further featurize each node with a unique relational context . We construct a relational context of a node n by randomly sampling a set of m immediate unique outgoing relations starting from n , i.e. , rconn = { rj } m ⊆ Nr ( n ) where Nr ( n ) denotes all outgoing relation types . Due to a non-uniform degree distribution , if |Nr ( n ) | < m , we add auxiliary [ PAD ] tokens to complete rconn to size m .

Summary: This paper presents NodePiece, a method inspired by subword embeddings in NLP that is designed for constructing compositional representations of entities in a knowledge graph using a fixed-size entity and relation vocabulary. This allows for learning and storing entity representations with a number of parameters that does not scale with the size of the knowledge graph, as well as being able to construct representations for unseen entities. The method uses a vocabulary that combines relation embeddings for all relations in the knowledge graph with a limited set of anchor entities, typically much fewer than the total number of entities. To construct an entity representation, a subset of that entity's nearest neighbor anchor entities (along with embeddings of the lengths of the shortest paths to those entities) are combined with a sampled set of adjacent relation embeddings and passed through an encoder (either an MLP or a Transformer) to output a fixed-size embedding. These compositional entity representations can then be used in classicial knowledge graph tasks such as link prediction, relation prediction, and node classification. Across all of these tasks and a range of datasets, NodePiece representations are shown to achieve a large fraction of the performance of strong baselines and in some cases even outperform them, all while using a fraction of their parameter count. Ablation analyses show that increasing the total number of anchors and number of anchors per entity improves performance up to a saturation point, and that in some cases a vocabulary of relations without anchor entities is sufficient (or even preferable) for constructing well-performing entity representations for certain downstream tasks.