N ODE P IECE : C OMPOSITIONAL AND P ARAMETERE FFICIENT R EPRESENTATIONS OF L ARGE K NOWLEDGE G RAPHS
A BSTRACT
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-efﬁcient node embedding strategies.
To this end, we propose NodePiece, an anchor-based approach to learn a ﬁxed-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 ﬁxed-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 classiﬁcation, 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 1 .
I NTRODUCTION
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 f n : n → R d and f r : r → R d which scale linearly to the number of nodes and edge types 2 , i.e., having O ( | N | ) space complexity of the entity embedding matrix.
Albeit efﬁcient 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 ×
200 d 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 inﬁnite combinations with a ﬁnite 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 ﬁxed-size token vocabularies, e.g., BERT (Devlin et al., 2019) contains ~30K tokens and
[Table]
GPT-2 (Radford et al., 2019) employs ~50K tokens.
Importantly, relatively small input embedding matrices enabled investing the parameters budget into more efﬁcient 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 ﬁxed 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 ﬁxed-size vocabulary V ( | V | (cid:28) | 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 ﬁrst 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 ) → R d which can be designed leveraging inductive biases of an underlying graph or downstream tasks.
Therefore, the overall parameter budget is now deﬁned by a small ﬁxed-size vocabulary of atoms and the complexity of the encoder function.
Our experimental ﬁndings suggest that a ﬁxed-size NodePiece vocabulary paired with a simple encoder still yields competitive results on a variety of tasks including link prediction, node classiﬁcation, 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.
R ELATED W ORK
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.
r r r r1_inv r3_inv a a a a a a
Closest anchors+
Anchor distances
Relational context  r1_inv r3_inv r
Encoder
[Figure]
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-speciﬁc 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.
N ODE P IECE V OCABULARY C ONSTRUCTION
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, ﬁxed-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 ) → R d 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. A NCHOR S ELECTION
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 justiﬁed when optimizing for certain task-speciﬁc topological characteristics, e.g., degree and PPR strategies indeed skew the distribution of shortest anchor distances towards smaller values thus increasing chances to ﬁnd anchors in 2- or 3-hop neighborhood of any node (we provide more evidence for that in Appendix C).
3. N ODE T OKENIZATION
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 (cid:0) | A | k (cid:1) 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 deﬁned by the total number of nodes, e.g., (cid:0)
|
A | k (cid:1) ≥ | N | .
Note that running depth-ﬁrst search (DFS) to random anchors at inference time is inefﬁcient 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-ﬁrst 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 distances 3 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 a i , we deﬁne anchor distance z a i ∈ [0; diameter ( G )] as an integer denoting the shortest path distance between a i and n in the original graph G .
Note that when tokenizing an anchor a j with the deterministic strategy, the nearest anchor among top- k is always a j itself with distance 0 .
We then map each integer to a learnable d -dimensional vector f z : z a i → R d akin to relative distance encoding scheme.
Relational Context.
We also leverage the multi-relational nature of an underlying KG.
Commonly 4 , the amount of unique edge types in G is orders of magnitude smaller than the total number of nodes, i.e., | R | (cid:28) | N | .
This fact allows to include the entire | R | in the NodePiece vocabulary V NP 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., rcon n = { r j } m ⊆
N r ( n ) where N r ( n ) denotes all outgoing relation types.
Due to a non-uniform degree distribution, if |N r ( n ) | < m , we add auxiliary [PAD] tokens to complete rcon n to size m .
3. E NCODING
At this step, a node n is tokenized into a sequence of k anchors, their k respective distances, and relational context of size m :
[Equation]
Taking anchors vectors a n and relation vectors r n from the learnable NodePiece vocabulary V ∈ R | V |× d , and anchor distances z a n from Z ∈ R ( diameter ( G )+1)
× d , we obtain a vectorized hash:
[Equation] [Equation]
As of 2021, one of the largest open KGs Wikidata contains about 100M nodes and 6K edge types
Although other operations are certainly possible, in this work, we use anchor distances as positional encodings of corresponding anchors and sum up their representations that helps to maintain the overall hash dimension of ( k + m )
× d .
Finally, an encoder function enc : R ( k + m ) × d → R d is applied to the vectorized hash to bootstrap an embedding of n .
In our experiments, we probe two basic encoders: 1) MLP that takes as input a concatenated hash vector R 1 × ( k + m )
d projecting it down to R d ; 2) Transformer encoder (Vaswani et al., 2017) with average pooling that takes as input an original sequence R ( k + m ) × d .
While MLP is faster and better scales to graphs with more edges, Transformer is slower but requires less trainable parameters.
As the two encoders were chosen to illustrate the general applicability of the whole approach, we leave a study of even more efﬁcient and effective encoders for future work.
While the nearest-neighbor hashing function has a greater number of collisions, its non-arbitrary mapping means that it is effectively permutation invariant.
We show this in Proposition 1 through the framework of Janossy pooling and permutation sampling based SGD, π -SGD (Murphy et al., 2019).
A proof is provided in Appendix H.
Proposition 1.
The nearest-anchor encoder with (cid:0) | A | k (cid:1) anchors and | m | subsampled relations, can be considered a π -SGD approximation of ( k + | m | ) -ary
Janossy pooling with a canonical ordering induced by the anchor distances.
Janossy pooling with π -SGD can be used to learn a permutation-invariant function from a broad class of permutation-sensitve functions such as MLPs (Murphy et al., 2019).
The permutation-invariant nature of the nearest-neighbor encoding scheme combined with the lack of transductive features such as node-speciﬁc embeddings mean that NodePiece can be used for inductive learning tasks as well.
With a ﬁxed-size vocabulary V NP , the overall complexity and parameter budget of downstream models are largely deﬁned by the complexity of the encoder and its inductive biases.
By design, the
NodePiece smaller vocabulary - larger encoder framework is similar to various Transformer-based language models (Qiu et al., 2020) whose vocabulary size remains rather stable with the encoder being the most important part responsible for the ﬁnal performance.
E XPERIMENTS
We design the experimental program not seeking to outperform the best existing approaches but to show the versatility of NodePiece on a variety of KG-related tasks: transductive, inductive, out-ofsample link prediction, and node classiﬁcation (with relation prediction results in Appendix I).
With this desiderata, we formulate the following research questions: RQ 1) Is it necessary to map each node to a unique vector for an acceptable performance on KG tasks?
; RQ 2)
What is the effect of hashing features?
; RQ 3) Is there an optimal number of anchors per node, after which diminishing returns hit the performance?
4. T RANSDUCTIVE L INK P REDICTION
Setup.
We run experiments on ﬁve KGs of different sizes (Appendix A.1) varying the total number of nodes from ~15K to ~2.5M. As a baseline, we compare to RotatE (Sun et al., 2019) that remains one of state-of-the-art shallow embedding models for transductive link prediction tasks.
To balance with NodePiece, RotatE operates on a graph with added inverse edges as well.
We report MRR with
Hits@10 in the ﬁltered (Bordes et al., 2013) setting as evaluation metrics, and count parameters for all models.
On larger KGs, we also compare to a smaller RotatE with a similar parameter budget.
In this task, NodePiece is equipped with a 2-layer MLP encoder.
For a fair comparison, we also adopt the RotatE scoring function as a link prediction decoder.
As to the NodePiece conﬁguration, we generally keep the number of anchors below 10% of total nodes in respective graphs.
We select 1k/ for FB15k-237 (i.e., total 1000 anchors and 20 anchors per tokenized node) with 15 unique outgoing relations in the relational context; 500/50 with 4 relations for WN18RR; 7k/20 with 6 relations for
CoDEx-L; 10k/20 with 5 relations for YAGO 3-10.
Other hyperparameters are listed in Appendix A.
Discussion.
Generally, the results suggest that a ﬁxed-size NodePiece vocabulary of <10% of nodes sustains 80-90% of Hits@10 compared to 10x larger best shallow models.
Some performance loss is expected due to the compositional and compressive nature of entity tokenization.
On smaller graphs
[Table] [Table]
(Table 2), parameter saving might not be well pronounced due to the overall small number of nodes to embed.
Still, taking even as few as 500 nodes as anchors on WN18RR retains 90% of the best model performance.
On bigger graphs (Table 3), parameter efﬁciency is more pronounced, i.e., on YAGO 3-10, a RotatE model of comparable size is 20 Hits@10 points worse than a NodePiece-based one.
This observation can be attributed to the fact the shrinking shallow models results in shrinking the embedding dimension of each node (20d for RotatE) which is inefﬁcient on small parameter budgets.
In contrast, a small ﬁxed-size vocabulary allows for larger anchor embedding dimensions (100d for
NodePiece with RotatE) since most of the parameter budget is deﬁned by the encoder.
We further study the effect of different anchor selection combinations (Fig. 2).
On WN18RR, fewer anchors with fewer anchors per node ( | A | /k ) yield relatively low accuracy but starting from
/(~0.1% of 40k nodes in the graph) the Hits@10 performance starts to saturate.
On FB15k-237, as few as 25 anchors already exhibit the signs of saturation where a further increase to 500 or 1000 anchors only marginally improves the performance.
We hypothesize such a difference can be explained by graph density, e.g., WN18RR is a sparse graph with a diameter of 23 and average anchor distance of 6 hops; while FB15k-237 is a denser graph with an average anchor distance of 2-3.
Hence, on a sparse graph with longer distances, it takes more anchors to properly encode a node.
However, more precise predictions (e.g., Hits@1) reﬂected in the MRR metric (see Appendix E) still remain a challenging task for small vocabulary NodePiece setups, and bigger | A | /k combinations alleviate this issue.
We also observe that diminishing returns, which make further vocabulary increase less rewarding, start to appear from anchor set sizes of ~1% of total nodes.
Ablations.
In the ablation study, we measure the impact of relational context and anchor distances on link prediction (Table 2).
Removing relational context and anchor distances does not tangibly affect the denser FB15k-237 data but does impair the accuracy on the sparser WN18RR.
Pushing vocabulary sizes to the limit, we also investigate NodePiece behavior in the absence of anchors at all, i.e., when hashes are deﬁned only by the relational context of size m .
Interestingly, this still yields fair performance on FB15k-237 with just 7 points Hits@10 drop, but drops to zero the WN18RR performance.
The fact that node embeddings might not be at all necessary but relations are more important supports the recent ﬁndings of Teru et al.
(2020) that relies only on relations seen in a small subgraph around a target node.
However, at this point, it seems to be a virtue of graphs with a diverse set of unique relations.
That is, FB15k-237 has 20x more unique relations than WN18RR and resulting hashes have more diverse combinations of relations which lead to more discriminative node representations.
Additionally, we visualize anchor embedding projections in Appendix D.
[Figure]
4.1. OGB W IKI KG
[Table]
To measure the beneﬁts of NodePiece on large-scale KGs
[Figure]
we run a link prediction experiment on OGB WikiKG
(Hu et al., 2020), a subset of Wikidata that consists of 2.5M nodes and 16M edges.
NodePiece is conﬁgured to sample a vocabulary 20K anchor nodes ( < 1% of total nodes) where each node is represented with k = 20 nearest anchors and a relational context of size m =
[Figure]
and we use a 2-layer MLP as a hash encoder (other hyperparameters as in Appendix A).
Generally, such a NodePiece conﬁguration can be paired with any link prediction decoder and we chose a non-parametric AutoSF (Zhang et al., 2020b) as one of the strongest decoders on this graph.
Overall, the NodePiece + AutoSF model has only 6.9M parameters, about 70 × smaller than top shallow models.
Compared to the best reported shallow approaches, the NodePiece-enabled model exhibits (cf.
Table 4, averaged over 10 seeds) even better performance achieved with a orders of magnitude smaller parameter budget.
We believe this result shows the effectiveness of NodePiece on large KGs with a dramatic parameter size reduction without signiﬁcant performance losses.
Ablations report the duality of a relational context, i.e., removing it from hashes leads to even higher MRR scores.
On the other hand, the relational context alone with 0 learnable anchors still yields considerably better results than 1000 × larger shallow models RotatE and TransE.
4. I NDUCTIVE L INK P REDICTION
We conduct a set of experiments on the inductive link prediction benchmark introduced by Teru et al.
(2020) to measure the performance of NodePiece features in the extreme case when anchor nodes are not available and only relational context can be used to compose entity representations.
Setup.
The unique feature of this benchmark compared to other evaluated tasks is that training and inference graphs are disjoint, i.e., inference at validation and test time is performed on a completely new graph comprised of new entities, and link prediction involves only entities unseen during training.
As inference graphs are disconnected from training ones, learning anchors from the training graph is useless, so node hashes are built only using the m -sized relational context.
On top of the obtained
NodePiece features we then employ a relational message passing GNN, CompGCN (Vashishth et al., 2020), with RotatE (Sun et al., 2019) as a scoring function for triples.
More details on the setup and best hyperparameters for NodePiece are presented in Appendix J.
Baselines .
We compare NodePiece + CompGCN with two families of models applicable in the inductive setting, i.e., rule-based methods, Neural LP (Yang et al., 2017), DRUM (Sadeghian et al., 2019), RuleN (Meilicke et al., 2018), and GNNs: GraIL (Teru et al., 2020) and recently proposed
Neural Bellman-Ford Nets (NBFNet) (Zhu et al., 2021).
Discussion.
Generally, the results conﬁrm the trend identiﬁed previously: relation-only features are strong performers in dense relation-rich graphs.
NodePiece features paired with CompGCN signiﬁcantly improve over path-based methods where performance gap might reach 37 absolute
Hits@10 points, e.g., in FB15k-237 V1 and NELL-995 V1.
Comparing to GNNs, NodePiece +
[Table] [Table]
CompGCN outperforms GraIL by a large margin in all (except one) experiments and is competitive to NBFNet on relation-rich FB15k-237 splits.
As expected, NodePiece features are less efﬁcient on sparse graphs (like WN18RR with few unique relations) but still outperform topology-based GraIL.
4. N ODE C LASSIFICATION
Setup.
Due to the lack of established node classiﬁcation datasets on multi-relational KGs, we design a multi-class multi-label task based on a triple version of a recent WD50K (Galkin et al., 2020) extracted from Wikidata.
The pre-processing steps are described in Appendix F
[Table]
and the ﬁnal graph consists of 46K nodes and 222K edges.
The task belongs to the family of transductive (the whole graph is seen during training)
semi-supervised (only a fraction of nodes are labeled) problems, where labels are 465 classes as seen in Wikidata.
In a semi-supervised mode, we test the models on a graph with 5% and 10% of labeled nodes.
Node features have to be learned as node embeddings.
As baselines, we compare to a 2-layer MLP and CompGCN (Vashishth et al., 2020) in a full-batch mode which is one of the strongest GNN encoders for multi-relational KGs.
Both baselines learn a full entity vocabulary.
We report ROC-AUC, PRC-AUC, and Hard Accuracy metrics as commonly done in standard graph benchmarks like OGB (Hu et al., 2020).
For PRC-AUC and Hard Accuracy, we binarize predicted logits using a threshold of 0.5.
Hard Accuracy corresponds to the exact match of a predicted sparse 465-dimensional vector to a sparse 465-dimensional labels vector.
NodePiece is conﬁgured to have only 50 anchors and use 10 nearest anchors per node with 5 unique relations in the relational context.
The dimensionality of anchors and relations is the same as in the baseline CompGCN.
Each epoch, we ﬁrst materialize all entity embeddings through the NodePiece encoder and then send the materialized matrix to CompGCN with the class predictor.
Discussion.
Surprisingly, ~1000x vocabulary reduction ratio (50 anchors against 46k for shallow models) greatly outperforms the baselines (Table 6).
MLP, as expected, is not able to cope with the task producing random predictions.
CompGCN, in turn, outperforms MLP demonstrating nonrandom outputs as seen by the ROC-AUC score of 0.836 and higher PRC-AUC and Hard Accuracy metrics.
Still, a NodePiece-equipped CompGCN with 50 anchors reaches even higher ROC-AUC of 0.98 with considerable improvements along other metrics, i.e., +16-19 PRC-AUC points and 3x boost along the hardest accuracy metric.
We attribute such a noticeable performance difference to better generalization capabilities of the NodePiece model.
That is, a generalization gap between training and validation metrics of the NodePiece + CompGCN is much smaller compared to the baselines who overﬁt rather heavily ( cf.
the training curves in the Appendix G).
The effect remains after increasing
[Table]
the number of labeled nodes to 10%.
Even with 50 anchors, the overall performance is saturated as the further increase of the vocabulary size did not bring any improvements.
Ablations.
We probe setups where NodePiece hashes use only anchors or only relational context
[Table]
and ﬁnd they both deliver a similar performance.
Following the previous experiments on dense graphs with lots of unique relations, it appears that node classiﬁcation can be performed rather accurately based only on the node relational context which is captured by NodePiece hashes.
4. O UT - OF - SAMPLE L INK P REDICTION
Setup.
In the out-of-sample setup, validation and test splits contain unseen entities that arrive with a few edges connected to the seen nodes.
For this experiment, we use the out-of-sample FB15k237 split (oFB15k-237) as designed in Albooyeh et al.
We do not employ their version of
WN18RR as the split contains too many disconnected entities and components in the train graph.
Instead, using the authors script, we sample a much bigger out-of-sample version of YAGO 3-10.
As a baseline, we compare to oDistMult (Albooyeh et al., 2020) which aggregates embeddings of all seen neighboring nodes around the unseen one (akin to 1-layer message passing with mean aggregator).
We adopt the same evaluation protocol - given an unseen node with its connecting edges, we mask one of the edges and predict its tail or head using the rest of the edges, repeating this procedure for each edge.
We report ﬁltered MRR and Hits@10 as main metrics.
NodePiece enables traditional transductive-only models to perform inductive inference as both seen and unseen nodes are tokenized using the same vocabulary.
For a smaller oFB15k-237 the NodePiece vocabulary has 1k/20 conﬁguration with 15 relations, while in a bigger oYAGO 3-10 we use 10k/ with 5 relations.
For this task, we apply a transformer encoder instead of MLP.
For a fair comparison, we use DistMult as a scoring function as well.
Discussion.
The results in Table 7 show that a simple NodePiece-based model retains ~90% of the baseline performance on oFB15k-237, but achieved faster and computationally inexpensive compared to oDistMult.
Moreover, while oDistMult is tailored speciﬁcally for the out-of-sample task, we did not do any task-speciﬁc modiﬁcations to the NodePiece-enabled model as it is inductive by design.
Furthermore, oDistMult is not able to scale to a bigger oYAGO 3-10 on a 256 GB RAM machine due to the out of memory crash.
Conversely, a NodePiece-equipped model has the same computational requirements as in other tasks and converges rather quickly (40 epochs).
Performed ablations underline the importance of having both anchors and relational context for tokenizing unseen entities.
We elaborate more on possible inference strategies for transductive and inductive tasks in Appendix A.6.
C ONCLUSION
In this paper, we have introduced NodePiece, a compositional approach for representing nodes in multi-relational graphs with a ﬁxed-size vocabulary.
Similar to subword units, NodePiece allows to tokenize every node as a combination of anchors and relations where the number of anchors can be 10–100 × smaller than the total number of nodes.
We show that in some tasks, node embeddings are not even necessary for getting an acceptable accuracy thanks to a rich set of relation types.
Moreover,
NodePiece is inductive by design and is able to tokenize unseen entities and perform downstream prediction tasks in the same fashion as on seen ones.
Reproducibility Statement.
The source code is openly available on GitHub.
All hyperparameters and implementation details are presented in Appendix A. Information on the used datasets is presented in Table 8 and we provide more details on dataset construction for node classiﬁcation and out-ofsample link prediction tasks in Appendix F.
The proof for Proposition 1 is given in the Appendix H.
Ethics Statement.
As NodePiece is a general graph representation learning method, we do not foresee immediate ethical consequences pertaining to the method itself.
Acknowledgements.
The authors would like to thank Koustuv Sinha, Gaurav Maheshwari and
Priyansh Trivedi for insightful and valuable discussions at earlier stages of this work.
We also thank anonymous reviewers for the helpful comments.
This work is partially supported by the Canada
CIFAR AI Chair Program and Samsung AI grant (held at Mila).
We thank Mila and Compute Canada for access to computational resources
A I MPLEMENTATION & H YPERPARAMETERS
[Table]
receive a new graph disjoint from the training one, and queries are sent against this new inference graph (hence the number of entities and facts for validation and test is the same)
[Table]
NodePiece is implemented in Python using igraph library (licensed under GNU GPL 2) for computing centrality measures and perform basic tokenization.
Downstream tasks employ NodePiece in conjunction with PyTorch (Paszke et al., 2019)
(BSD-style license), PyKEEN (Ali et al., 2021) (MIT License), and PyTorch-Geometric (Fey & Lenssen, 2019)
(MIT License).
We ran experiments on a machine with one RTX 8000 GPU and 64 GB RAM.
The OGB WikiKG 2 experiments were executed on a single Tesla V100 16 GB VRAM and 64 GB RAM.
All used datasets are available under open licenses.
For all downstream tasks and datasets we employ the deterministic anchor selection strategy where 40% of the total number of anchors | A | are nodes with top PPR scores, 40% are top degree nodes, and remaining 20% are selected randomly.
All anchor sets are non-overlapping and disjoint, i.e., if some top degree nodes have already been selected with the PPR policy, they will be skipped in favor of next nodes in the sorted list.
The choice for this strategy is motivated in Appendix C.
A. D ATASETS
Details on the datasets for transductive link prediction, out-of-sample link prediction, relation prediction and node classiﬁcation are collected in Table 8.
The inductive link prediction benchmark introduced by Teru et al.
(2020) includes 3 graphs, FB15k-237, WN18RR, and NELL-995, each hasdifferent splits that vary in the number of unique relations, number of nodes and triples at training and inference time.
Full dataset statistics is provided in Table 9.
FB15k-237 and most splits of NELLcan be considered as relation-rich graphs while WN18RR is a sparse graph with few relation types.
A. T RANSDUCTIVE L INK P REDICTION
The optimizer is Adam for all experiments.
As RotatE is a scoring function in the complex space, the reported embedding dimensions are a sum of real and imaginary dimensions, e.g., 1000d means that both real and imaginary vectors are 500d.
[Table]
[Table]
A. R ELATION P REDICTION
Conﬁgurations (Table 12) for the compared models are almost identical to those of the transductive link prediction experiment.
We mostly reduce the number of epochs and negative samples as models converge faster on this task.
A. N ODE C LASSIFICATION
In this experiment (Table 13), NodePiece is used at the initial step to bootstrap a node embeddings matrix which is then sent to the CompGCN graph encoder.
In contrast, CompGCN and MLP baselines use directly a trained node embedding matrix as their initial input.
A. O UT - OF - SAMPLE L INK P REDICTION
The set of NodePiece hyperparameters (Table 14) is similar to the set of the transductive experiments except the scoring function (DistMult), encoder function (Transformer), and number of epochs as the model converges faster.
We do not provide a setup for the baseline oDistMult on oYAGO 3-10 as the model was not able to pre-process the dataset on a machine with 256 GB RAM.
Reported training
[Table] [Table]
times for NodePiece models exclude evaluation.
Training times of the baseline oDistMult were not reported by its authors.
A. D EPLOYMENT IN R EAL - WORLD D YNAMIC K NOWLEDGE G RAPHS
NodePiece, on account of its compositional representation, can be applied to dynamic real-world knowledge graphs where nodes are added and removed over time.
That is, training on a graph snapshot we can obtain the embeddings of new nodes without re-computing and updating representations of every other node in the graph.
This is valuable in settings where there are latency requirements such as many online services.
For example, if a user creates an account on a social media service and begins liking content (represented as a “like” edge in the social network graph between the user and the content), it would be desirable to make future content recommendations rapidly reﬂect this new data without waiting for the next batched retraining to update the users embedding.
In order to reduce latency the entire embedding matrix can be materialized and cached ahead of time, updating embeddings as new nodes and edges are added.
The parameter efﬁciency and compositionality of
NodePeice means that for large real-world graphs NodePiece subsumes what would before have been a complex system of a large-scale embedding framework like Pytorch-BigGraph (Lerer et al., 2019), an OOV embedding method (e.g,.
ERAvg) and a shallow embedding method (e.g., RotatE).
[Table]
B L IMITATIONS AND F UTURE W ORK
Our experimental results demonstrate the promise of using NodePeice to signiﬁcantly reduce the parameter complexity of node embeddings.
While it is difﬁcult to prove
[Table]
we also hypothesize that the parameters required by NodePeice to maintain the same level of performance (as the graph scales) increase sublinearly according to the size of the graph.
The intuition for this is twofold.
First, the number of unique anchor combinations of size k that can be encoded increases according to (cid:0) | A | k (cid:1) (i.e. O ( | A | k ) ) if randomly sampled — if the sampling is done via nearest neighbor anchor selection then the number of unique permutations is expected to increase polynomially.
Second, increasing the size of the graph will only require sublinear increase in the number of anchors in order to maintain the same average node-anchor distance.
Although proving causality is difﬁcult, we believe that maintaining hashing uniqueness and node-anchor distances stable will be sufﬁcient to maintain equivalent performance.
C A NCHOR S ELECTION S TRATEGIES
Here, we provide more details as to anchor conﬁgurations ( k nearest from total A anchors) and anchor distances.
Recall that there exist several ways to select the total set of anchors
A as stated in Section 3.1, i.e., random or centrality-measure based.
Then, k anchors per node can be chosen either as k nearest (default NodePiece mode) or k random anchors.
Figure 3 depicts the effect of those strategies on the distribution of anchor distances (number of hops between a target node and its anchors).
We use the conﬁgurations used in the main experiments, i.e., 1000 anchors and 20 anchors per node for FB15k-237, and 500 anchors with 50 anchors per node on WN18RR.
First, we observe that PPR, degree, and mixed (40% PPR, 40% degree, 20% random) strategies generally skew the distribution towards smaller anchor distances compared to random strategies.
This fact supports the hypothesis that deterministic strategies improve the chances to ﬁnd an anchor in a closer l -hop neighborhood of a target node.
Second, varying the way of selecting k anchors per node between nearest (left column) and random (right column), we also observe the skew of a distribution of anchor distances.
Next, we ﬁx the anchor selection strategy to the mix , ﬁx the number of anchors per node (50 for
WN18RR and 20 for FB15k-237), and vary a total number of anchors A (50 to 1000 for WN18RR and 20 to 1000 for FB15k-237) along with the method of sampling k anchors per node, i.e., nearest and random.
Figure 4 shows that increasing the total number of anchors together with k nearest
[Figure] [Figure]
anchors again skews the distribution of anchor distances towards smaller values and, hence, to higher probabilities of ﬁnding anchors in a closer neighborhood of a target node.
We would recommend using centrality-based strategies to select A with k nearest anchors per node if anchor distances and probability of ﬁnding anchors in a closer neighborhood are of higher importance.
Finally, we ﬁx the anchor selection strategy as mix , obtain nearest anchors per node, and under this setup study average anchor distances varying k - the number of anchors per node in various combinations of total anchors A .
The results presented on Figure 5 suggest that sparser graphs (like
WN18RR) beneﬁt more from increasing the number of anchors A , i.e., the delta between distances is much larger than that of dense FB15k-237.
The difference in distances on Figure 5 might also explain the performance on Figure 8, i.e., generally, smaller A/k conﬁgurations like 25/5 are inferior on sparser graphs but perform competitively on denser ones.
[Figure]
D E MBEDDING V ISUALIZATIONS
To further study learned representations of anchors and capabilities of the encoder, we build tSNE and UMAP projections from subsamples of FB15k-237 and WN18RR based on trained models from the transductive link prediction experiments (hyperparameters listed in Table A).
For FB15k-237, we randomly sample 1000 entities (out of total 15K) and ﬁnd their top-100 most common anchors.
The anchor embeddings are extracted from the learned tensor while 1000 entity embeddings are obtained through the NodePiece encoder.
Similarly for WN18RR, we sampleentities (out of total 40K) keeping their top-100 most common anchors.
As we use the RotatE decoder that assumes entities and anchors are modeled in a complex space, we visualize their real parts (e.g., ﬁrst 100 dimensions out of 200).
Recalling that link prediction performance of NodePiece + RotatE retains 80-90% of the state of the art models performance, the results on Fig.
6 and Fig.
7 demonstrate that (1) NodePiece encoder is able to reconstruct clusters of similar entities; (2) anchors are well-scattered among communities.
Albeit entity embeddings are built as a composition of k anchors, it can be seen that all communities have "specialized" nearby anchors.
On a higher level, common anchors tend to be well-scattered in the space.
Less common anchors, as seen on FB15k-237 and Fig. 6, tend to group together.
However, thanks to the non-linear nature of the NodePiece encoder, resulting entity embeddings still form different clusters and communities not concentrated around one point.
We believe this is the effect of a compositional encoder and plan to investigate this phenomenon further.
E T RANSDUCTIVE L INK P REDICTION R ESULTS : MRR
In addition to Figure 2 that presents Hits@10, we report variations of mean reciprocal rank (MRR) depending on combinations of A and k on Figure 8 from the same set of experiments.
On sparser
WN18RR, smaller A/k combinations like 25/5 or 50/10 struggle with more precise predictions like
Hits@1 which is captured by low values of MRR.
Starting from 500/10, the WN18RR performance starts to saturate.
On the other hand, on denser FB15k-237, the difference between minimum and maximum MRR is less than 4 points, and performance exhibits signs of saturation already at 50/10.
F D ATASETS C ONSTRUCTION F. N ODE C LASSIFICATION : WD50K NC
The original WD50K (Galkin et al., 2020) contains a triple-only KG version on which we base a new dataset for semi-supervised multi-class multi-label node classiﬁcation.
First, we remove all triples containing Wikidata properties P31 ( instance of ) and P279 ( subclass of ) as they already contain class information.
We then remove nodes that became disconnected after removing those edges.
Third, using SPARQL queries, for each remaining node in a graph, we extract a 3-hop class hierarchy of Wikidata classes and their superclasses.
We only keep class labels that occur at least 50 times in the training set.
Then, we sample 10% of nodes with labels for validation and 10% for test, and of remaining 80% we sample a set of nodes for the semi-supervised setup, i.e., we keep only 5% and
[Figure]
10% of those nodes.
The resulting graph has 46k nodes, 526 distinct relation types, and 465 class labels.
F. O UF - OF - SAMPLE L INK P REDICTION : O YAGO 3-
For sampling the out-of-sample version of a bigger YAGO 3-10 we largely follow the same original procedure described in Section 4 of (Albooyeh et al., 2020).
We ﬁrst merge the train, validation and test triples from the original dataset for transductive link prediction.
Then, from all entities appearing in at least two triples, we randomly sample 5% of nodes to be the out-of-sample entities for validation and 5% for test.
All triples containing the out-of-sample entities on subject or object positions are put into validation or test, respectively, as edges that connect an unseen entity with the seen graph.
[Equation]
Figure 9 depicts train and validation values of Hard Accuracy and PRC-AUC metrics for all the compared models on WD50K NC with 5% of labeled nodes.
The NodePiece model has only 50 total anchors with 10 nearest anchors per node, and 5 unique relation types in the relational context.
The performance on the dataset with 10% of nodes is almost the same, so we report the charts only on 5% dataset.
By the generalization gap we understand the delta between training and validation values.
The MLP baseline quickly overﬁts but fails to generalize on the validation.
The generalization gap of CompGCN is smaller compared to MLP but is still signiﬁcant, i.e., validation performance is 2–3 × smaller than train.
Finally, the NodePiece-enabled model has the smallest generalization gaps, especially along the Hard Accuracy metric where the validation performance is very close to that of train.
Similarly, the gap on PRC-AUC is smaller than 10 points.
As shown in the ablation study in Table 6, it appears that explicit node embeddings do not contribute to the classiﬁcation performance.
Hence, the baseline models tend to be overparameterized where learnable node embeddings add noise, while the NodePiece model has only a few anchors (or no anchors at all when using only the relational context), much fewer parameters, and therefore generalizes better.
This hypothesis also explains the observation that the node classiﬁcation performance does not improve when increasing A/k anchor conﬁgurations.
H P ROOFS
Proposition 2.
The nearest-anchor encoder with (cid:0) | A | k (cid:1) anchors and | m | subsampled relations, can be considered a π -SGD approximation of ( k + | m | ) -ary
Janossy pooling with a canonical ordering induced by the anchor distances.
Proof .
We begin by providing the deﬁnition of Janossy pooling as it was presented in the original paper
Murphy et al.
Deﬁnition 1 (Janossy pooling) .
Let H ∪ be the union of all anchors and relations.
Consider a function f (cid:42) :
N × H ∪ × R d → F on variable-length but ﬁnite sequences h , parameterized by θ ( f ) ∈ R d , d > 0 .
A permutation-invariant function f : N × H ∪ × R d → F is the Janossy function associated with f (cid:42) if
[Equation]
where Π | h | is the set of all permutations of the integers 1 to | h | , and h π represents a particular reordering of the elements of sequence h according to π ∈ Π | h | .
We refer the operation used to construct f from f (cid:42) as Janossy pooling.
While Janossy pooling provides a simple approach to construct permutation-invariant functions from arbitrary permutation sensitive functions, it is computationally intractable due to the need to sum over all computations.
Three general strategies proposed under this framework to overcome this combinatorial challenge: canonical orderings, k -ary
Janossy pooling, and π -SGD approximations.
A very effective way of reducing the complexity is to constrain the permutations to a canonical ordering that is independent of a speciﬁc adjacency matrix ordering over a given graph.
More precisely, one deﬁnes as a function CANONICAL : H ∪ → H ∪ such that CANONICAL( h ) =
CANONICAL( h π ) ∀ π ∈ Π | h | and only considers functions f (cid:42) based on the composition f (cid:42) =
CANONICAL ◦ f (cid:42) (cid:48)
(Murphy et al., 2019).
In the case of NodePiece we are able to deﬁne this ordering for the anchors according to their distance to the target node.
Assuming that the number of relations is ﬁxed or grows at slow rate throughout the life-cycle of a graph we can deﬁne an arbitrary ordering for relations as a canonical ordering for the relational context.
However, since anchors can be equidistant such a canonical ordering does fully satisfy permutation invariance.
We propose a trivial relaxation of the original deﬁnition of canonical orderings simply requiring that an ordering greatly reduce the number of unique permutations since in practice an exact canonical ordering is rarely feasible.
Speciﬁcally, |{ CANONICAL( h π ) ∀ π ∈ Π | h | }| (cid:28) |{ ( h π ) ∀ π ∈ Π | h | }| .
To further reduce the number of permutations we can truncate our ordered sequence h .
This is known as k -ary
Janossy pooling pooling (Deﬁnition 2) and is implicitly performed by the NodePeice algorithm by varying the anchor per node parameter, k , and the size of the relational context, | m | .
Deﬁnition 2 ( k -ary Janossy pooling) .
Fix k ∈ N .
For any sequence h , deﬁne ↓ k ( h ) as its projection to a length k sequence; in particular, if | h | ≥ k , we keep the ﬁrst k elements.
Then, a k -ary
[Table]
permutation-invariant Janossy function f is given by
[Equation]
Since an imperfect truncated canonical ordering may still result in a potentially intractable number of permutations
[Table]
we use permutation sampling also known as π -SGD to learn arbitrary functions that approximate ( k + | m | ) -ary
Janossy pooling.
This is done by randomly ordering anchors that are equidistant resulting in a uniform sampling of possible permutations during training and evaluation.
For more details on the formal deﬁnition of π -SGD we point the reader to the original paper (Murphy et al., 2019).
I R ELATION P REDICTION
Setup.
We conduct the relation prediction experiment on the same FB15k-237, WN18RR, and YAGO 3-10 datasets.
While link prediction deals with entities, the relation prediction model has to rank a correct relation given a (head, ?, tail) query.
We report MRR and Hits@10 in the ﬁltered setting as evaluation metrics.
Similar to the link prediction conﬁguration, we use NodePiece + 2-layer MLP and compare against RotatE of the same total parameter count.
Discussion.
The reported results (Table 15) demonstrate a competitive performance of NodePiecebased models with reduced vocabulary sizes bringing more than 97% Hits@10 across graphs of different sizes.
In the case of WN18RR and YAGO 3-10, NodePiece models with fewer anchors even slightly improve the accuracy upon the shallow embedding baseline.
The ablation study suggests that on dense graphs with a reasonable amount of unique relations having explicit learnable node embeddings might not be needed at all for this task.
That is, we see that on FB15k-237 and YAGO 3-10 the NodePiece hashes comprised only of the relational context deliver the same performance without any performance drop conﬁrming the ﬁndings from the previous experiment.
J I NDUCTIVE L INK P REDICTION
Setup .
Nodes in the inference graphs do not have any associated feature vectors which makes this benchmark very relevant for graph representation learning.
Importantly, the set of relation types in the inference graphs is a subset of those seen in the training set.
Since the relation embedding matrix can be learned on the training graph, we therefore have a uniform method for constructing node representations both on seen and unseen graphs.
As an encoder we try MLP and Transformer.
Evaluation Protocol .
Following the original work (Teru et al., 2020), we employ a ﬁltered setting and rank each triple against 50 random negative triples reporting the Hits@10 metric.
This setup is motivated by computational complexity of GraIL at inference time, while NodePiece + CompGCN is as fast in the inductive inference as in the transductive regime reported in other experiments.
Discussion .
For each KG, there are 4 splits of increasing size of train and inference nodes and edges.
The empirical results on this spectrum of various sizes demonstrates interesting scalability properties of NodePiece in inductive settings.
Without anchor nodes, the NodePiece vocabulary size is independent of the number of nodes and edges, depending only on the number of relation types.
On dense relation-rich graphs such a vocabulary is enough to yield very competitive performance.
[Table]
K U NIQUENESS OF N ODE H ASHES
NodePiece represents nodes as a sequence of tokens and a natural question in this context is how unique such sequences can be in light of different anchor selection and tokenization strategies.
Assuming the input graph is a single connected component, when sampling | A | total anchors and selecting randomly k anchors per node
[Table]
the number of possible hash combinations is bounded by (cid:0) | A | k (cid:1) .
In this scenario, uniqueness of hashes is achieved by having this number bigger than the number
N of nodes in a graph, (cid:0) |
A | k (cid:1) > N , and this happens with high probability for any reasonably large | A | , eg, (cid:0) 5020 (cid:1) encodes about 4 .
7 · 10 13 combinations which covers all existing public KGs combined.
In the deterministic selection of nearest anchors per node we do not have such guarantees.
Nevertheless, additional sequences of relational contexts and anchor distances help to obtain more unique hashes.
Collisions are possible in highly regular graphs like Wordnet, but real-world KGs like
Wikidata and DBpedia do not exhibit such a regular structure.
Similar to homonyms whose meaning depends on the surrounding context in a sentence, we hypothesize that adding message passing layers (that can be seen as encoding of a neighboring context for a given node) on top of NodePiece hashes might further improve the diversity of node representations.
That said, the future work research agenda might include proving tighter theoretical bounds on hashes uniqueness and developing new anchor sampling and tokenization strategies.