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 to scale up GNNs by means of removing their dependence on individual node embeddings which grow linearly in size and are inefficient in huge graphs. The method uses a set of anchor nodes, which are picked from the graph itself, for the NodePiece representations. Concretely, NodePiece embedding is derived by the embeddings of the closest anchor nodes and the hop-based distance to these nodes which are also represented by vectors. – the maximum distance between two nodes in a graph is Diameter(G), so the number of distance/position embeddings is Diameter(G). Additionally, NodePiece is augmented by a sample of embeddings of the relationship types the node is involved in. Apart from induced efficiency, NodePiece can be especially useful in inductive learning where an unseen node’s embedding can be created through the anchor nodes.