SlowGuess/NAACL_2022_OCR huggingface dataset:
Elan Markowitz $^{1,2,\dagger,*}$ , Keshav Balasubramanian $^{1,*}$ , Mehrnoosh Mirtaheri $^{1,2}$ , Murali Annavaram $^{1}$ , Aram Galstyan $^{1,2}$ , and Greg Ver Steeg $^{1,2}$
<sup>1</sup>University of Southern California  
<sup>2</sup>USC Information Sciences Institute  
*Denotes equal contribution
†Corresponding author: esmarkow@usc.edu
Knowledge graphs (KGs) often represent knowledge bases that are incomplete. Machine learning models can alleviate this by helping automate graph completion. Recently, there has been growing interest in completing knowledge bases that are dynamic, where previously unseen entities may be added to the KG with many missing links. In this paper, we present StATIK-Structure And Text for Inductive Knowledge Completion. StATIK uses Language Models to extract the semantic information from text descriptions, while using Message Passing Neural Networks to capture the structural information. StATIK achieves state of the art results on three challenging inductive baselines. We further analyze our hybrid model through detailed ablation studies.
Knowledge graphs (KGs) are appropriate representations of knowledge bases across many domains. These domains include commonsense reasoning (Bauer, 2021; Yan et al., 2021; Zhang et al., 2020), question answering (Yasunaga et al., 2021; Feng et al., 2020; Lin et al., 2019; Christmann et al., 2019; Saxena et al., 2021; Hixon et al., 2015), recommendation systems (Guo et al., 2020; Wang and Cai, 2020; Huang et al., 2018; Wang et al., 2018), and many others (Hogan et al., 2021). These graphs are extremely large and often incomplete. As a result, there is significant interest in training machine learning models that can help complete these knowledge bases. In knowledge graphs, the nodes, called entities, often possess textual descriptions, while edges are typically labeled with one of many relation types, which may also possess textual descriptions. Effective KG completion models should learn to leverage this textual information in order to correctly complete the knowledge base. Additionally, such knowledge graphs are usually dynamic (Das et al., 2018; Liao et al., 2021) as
Figure 1: Depiction of the problem addressed by inductive learning. During training, only the blue portion of the graph exists, including the entities FC Barcelona, sports team, La Liga, and Spain. Later, the entity Cádiz FC is added to the graph. When added, an entity contains a description and some number of edges (possibly zero). Since StATIK is inductive, it requires no retraining or retroactive processing in any way to make predictions about Cádiz FC. This could include predicting Cádiz FC's country i.e. the query (Cádiz FC, country, ?). The correct prediction, (Cádiz FC, country, Spain), is displayed in dashed red.
a result of the underlying knowledge base being dynamic. Typically, this manifests as nodes and edges being added and removed from the knowledge graph while the set of relation types remains more static. Thus, another quality we desire of knowledge graph completion models, is that they be inductive and generalize to unseen entities.
We propose a completely inductive, hybrid model, called StATIK, that effectively leverages both the structure of a knowledge graph as well as the underlying textual descriptions of the entities and relations. Structure is incorporated through a Message Passing Neural Network (MPNN) (Gilmer et al., 2017) that aggregates information from a neighborhood defined around each entity, while textual information is incorporated through a pretrained language model such as BERT (Devlin et al., 2019). Our main contributions are summarized as follows:
1. We propose a completely inductive and hy
Findings of the Association for Computational Linguistics: NAACL 2022, pages 604 - 615 July 10-15, 2022 ©2022 Association for Computational Linguistics
Model	TransE	OpenWorld	Glove-DKRL	Commonsense	IndTransE	LAN	GraIL	KGBert	BLP	StAR	ours
Inductive - Seen2Unseen	X	✓	✓	X	✓	✓	✓	✓	✓	✓	✓
Inductive - Unseen2Unseen	X	✓	✓	X	X	X	✓	✓	✓	✓	✓
End-to-end LM	X	X	X	X†	X	X	X	✓	✓	✓	✓
No Support Set Required	✓	✓	✓	✓	X	X	X	✓	✓	✓	✓
Graph features	X	X	X	✓	✓	✓	✓	X	X	X	✓
Structure Objective	✓	✓	✓	✓	✓	✓	✓	✓	✓	X	✓
Inference Scalability	O(N)	O(N)	O(N)	O(N)	O(N)	O(N)	O(NQ)	O(NQ)	O(N)	O(NR+Q)	O(N+Q)
Table 1: Related works comparison table.  $N$  is number of entities,  $Q$  is number of queries,  $R$  is number of relation types.  $\dagger$ Model uses domain adaptation but does not train end-to-end. References are TransE (Bordes et al., 2013), OpenWorld(Shah et al., 2019), Glove-DKRL(Xie et al., 2016), Commonsense(Malaviya et al., 2020), IndTransE(Dai et al., 2021), LAN(Wang et al., 2019a), GraIL(Teru et al., 2020), KGBert(Yao et al., 2019), BLP(Daza et al., 2021), StAR(Wang et al., 2021)
brid knowledge graph completion model composed of an MPNN to leverage structure and a language model to leverage text.
2. We demonstrate empirically that incorporating structure via an MPNN leads to much better generalization in the inductive setting.
3. We achieve state-of-the-art results on three inductive benchmarks in which each predicted triple contains at least one new entity (dynamic setting) or exclusively new entities (transfer setting).
4. We design our model with scale in mind and show that the proposed model is significantly faster than similar alternatives, particularly for inference.
2 Inductive Representation Learning on Knowledge Graphs
We can define a knowledge graph with textual information as  $\mathcal{G} = (\mathcal{E},\mathcal{R},\mathcal{T},\mathcal{D})$  where  $\mathcal{E}$  is the set of entities,  $\mathcal{R}$  is the set of relation types,  $\mathcal{T}$  is the set of triples  $(h,r,t)\in \mathcal{E}\times \mathcal{R}\times \mathcal{E}$ , and  $\mathcal{D}$  is the set of entity and relation descriptions. The inductive graph completion task is defined as follows. Let the training graph be  $\mathcal{G}_{train} = (\mathcal{E}_{train},\mathcal{R},\mathcal{T}_{train},\mathcal{D}_{train})$  where  $\mathcal{E}_{train}$  is a subset of  $\mathcal{E}$ ,  $\mathcal{D}_{train}$  is the corresponding subset of  $\mathcal{D}$ , and  $\mathcal{T}_{train}$  is the subset of  $\mathcal{T}$  containing triples only involving entities in  $\mathcal{E}_{train}$ . The inference task is then to predict the triples in  $\mathcal{T} / \mathcal{T}_{train}$ , only having trained the model on  $\mathcal{G}_{train}$ . Specifically, for a given evaluation triple,  $\mathcal{T}_i = (h,r,t)$ , do head and tail prediction on the
graph  $\mathcal{G} - \mathcal{T}_i$ . This means that given a query of the form  $(h,r,?)$  or  $(?,r,t)$ , rank all possible tail or head candidates (targets) such that the real triple is ranked as highly as possible. Figure 1 demonstrates a motivating example.
Similar to the transductive setting, such a task can be solved by learning a model that scores triples through minimizing some objective. We use a margin ranking loss as our objective where given a set of real triples,  $T$ , a corresponding set of negative triples,  $T'$ , and a scoring function  $f$  (higher scores imply more likely triple), we compute the loss as
$$
\sum_ {(t, t ^ {\prime}) \in (T \times T ^ {\prime})} \max  \left(0, 1 - f (t) + f \left(t ^ {\prime}\right)\right) \tag {1}
$$
In order to learn the inductive objective, the model should avoid entity specific parameters such as an entity embedding tables as those parameters will not translate to the new entities at test time. To tackle this challenge StatIK uses text features instead of embeddings and extends prior work by also incorporating structural information through message passing.
Much of the work in the area of Knowledge Graph Completion has focused on the transductive setting i.e. performing link prediction between entities seen at training time. Generally, these methods learn embeddings in a geometric space such as TransE (Bordes et al., 2013), ComplEx (Trouillon et al., 2016), DistMult (Yang et al., 2015), RotatE (Sun et al., 2019), and SimplE (Kazemi and Poole,
(a) Model Flowchart
(b) Message Passing Layer
Figure 2: Depiction of model flowchart (left) and the MPNN component (right). Green indicates model inputs. Yellow indicates pre-processing steps. Purple indicates data loading steps. Blue indicates model computation steps. Solid lines represent inputs and outputs. Dashed lines indicate residual connections.
2018), or through a machine learning decoder such as ConvE and HypER (Dettmers et al., 2017; Balazevic et al., 2018). There has also been effort in using graph neural networks for knowledge graph completion. R-GCN (Schlichtkrull et al., 2018) brings the original GCN (Kipf and Welling, 2017) to the multi-relational knowledge graph setting. Wang et al. (2020b) looked at using a modified version of GAT (Velickovic et al., 2018) to get strong results in the transductive setting.
3.1 Inductivity
Recently, there has been increased focus on the inductive setting. Works such as LAN (Wang et al., 2019a) and IndTransE (Dai et al., 2021) as well as a few others look at learning embeddings for new entities based on edges to entities in the training graph (Wang et al., 2020a; Bhowmik and de Melo, 2020). This requires a sufficient number of edges from nodes seen during training to the new nodes (seen-to-unseen). Other methods, such as GraiL (Teru et al., 2020), have been able to achieve inductivity without such requirements, and as a result, can operate on unseen-to-unseen entities.<sup>1</sup> There have also been some works on open domain knowledge graph completion, a similar learning task. These works include Shah et al. (2019) (OpenWorld) and Shi and Weninger (2018), and some of their techniques, such as using text to enable generalization to new entities, have continued in the works analyzed here. OpenWorld specifically aims to learn a
function that aligns unsupervised text embeddings with knowledge graph embeddings so that new entities can be placed in the KG embedding space.
3.2 Support Sets
Most of the aforementioned inductive works (Wang et al., 2019a; Dai et al., 2021; Wang et al., 2020a; Bhowmik and de Melo, 2020) all require a support set, or edges connecting to known entities in the training graph, for the new entities seen at test time. While it is certainly useful to be able to use a support set when available, an ideal model would have the flexibility to use such edges when present but still be able to produce meaningful representations without them.
3.3 Language Models
Transformers (Vaswani et al., 2017), have created a renaissance in language modeling over recent years. Combined with self-supervised pretraining, language models such as BERT (Devlin et al., 2019) are able to capture the contextual and semantic information of natural language.
As many KGs contain text associated with each entity, researchers have sought to use that information for improved performance or inductivity. KGBert (Yao et al., 2019) looked into using transformers for link prediction, treating it as a text classification task. Bert for Link Prediction (BLP) (Daza et al., 2021) and StAR (Wang et al., 2021) have sought to incorporate language models while improving on some of the flaws of KGBert. Commonsense (Malaviya et al., 2020) also incorporates
${}^{1}$  GraiL is technically designed to do relation prediction, not link prediction, but is included here due to its relevance.
a language model (along with a Graph Neural Network). However, it only uses the language model to initialize an embedding table. As a result, it is not inductive like the other models. Older model DKRL (Xie et al., 2016) uses a simpler language model with GloVe embeddings (Pennington et al., 2014).
3.4 Structural Objective
Most KG completion models use some form of structural objective; The scoring function uses spatial or geometric transformations to capture the graph structure. For instance, TransE applies the structural objective that a head embedding  $+$  a relation embedding should be close to the tail embedding of a true triple. KGBert is one of the few models that does not use such an objective.
3.5 Graph features
Structural objectives alone have some limitations with regard to capturing graph structure. Being able to explicitly use the local graph structure and topology as a feature (often through message passing) is beneficial for both general performance and inductivity. Many of the models mentioned (Malaviya et al., 2020; Dai et al., 2021; Wang et al., 2019a; Teru et al., 2020) make use of such features.
3.6 Scalability
Scalability is an incredibly important aspect of KG completion as knowledge graphs can include millions to billions of entities and edges. When dealing with complex encoders such as MPNNs or LMs, the number of encoder passes becomes an especially pressing issue, even for smaller graphs.
In KG completion, every entity is considered a possible solution to a query. If each of these possible triples is evaluated independently, the problem quickly becomes a combinatorial mess. This is the case for KGBert and GraiL which can only evaluate a single triple at a time. This makes the link prediction task quadratic in complexity. StAR improves on KGBert's approach but requires evaluating each entity in conjunction with every relation type. While definitely an improvement, this can still be problematic for big graphs.
Table 1 gives a comparison of the most relevant related works.
4 StATIK Architecture
We can conceptualize KG completion models as having an encoder and a decoder. Whereas most
KG models use an embedding table for the encoder, our model utilizes a hybrid language model and MPNN based network. Our decoder is a relatively simple scoring function.
4.1 Language Models for Feature Extraction
For any entity or relation  $x$ , let  $L(x) \in \mathbb{R}^{d_0}$  be the  $d_0$  dimensional BERT-base encoding for the text associated with  $x$ . For every entity and relation type, we preprocess  $L(x)$  for all  $x$  to create a meaningful feature matrix.
We explore combining two different encoders, a transformer for processing text and a message passing neural network for processing structural data.
4.2.1 Language Model Transformer
Separately from the text featurization process, we also train through a language transformer in an end-to-end manner.
When encoding target entities, we simply pass the text associated with the entity through the encoder. However, when encoding queries, we condition the entity on the relation type. Similarly to StAR, for tail prediction-  $(h,r,\ref{eq:1})$  we append the text associated with the head and the relation together giving head_text  $^+$  relation_text. For queries of the form  $(?,r,t)$ , we prepend the relation text with the text "inverse of". This gives us the text for the query as tail_text + "inverse of"  $^+$  relation_text. In contrast, StAR only ever appends head entities to relation text. This makes the head prediction task more complicated as one would have to encode each possible head-relation pair as the target of the head prediction queries.
4.2.2 Message Passing Graph Neural Network (MPNN)
We employ an MPNN of the following form (see fig 2b for overview). For each query of the form  $(h,r,?)$  or  $(?,r,t)$ , let  $s$  be the "query entity", meaning the head or tail that is part of the query (or the target entity if encoding target candidates). We then form a subgraph around  $s$  using the edges,  $E$ , connecting  $s$  to its neighbors,  $\mathcal{N}(s)$ .
We then compute the initial representations of all the entities in the subgraph,  $V = \mathcal{N}(s) \cup \{s\}$ , as
$$
X _ {0} = \sigma (L (V) \mathbf {W} _ {\mathbf {0}} ^ {(\mathbf {e})}) \tag {2}
$$
where  $\mathbf{W}_0^{(\mathbf{e})}$  is a  $\mathbb{R}^{d_0\times d}$  parameter to reduce the dimension to the model's hidden dimension size  $d$ ;  $\sigma$  is the model's element-wise activation function (LeakyReLU in our case).
For the  $m$  edges in the subgraph, let  $\mathbf{i_h}$  and  $\mathbf{i_t}$  be  $\mathcal{E}^m$  vectors indicating the heads and tails respectively of each edge. Let  $\mathbf{i_r} \in \mathcal{R}^m$  be the relation type of each edge. Let  $\mathbf{i}_{\mathrm{dir}} \in \{0,1\}^m$  be the relative direction of the edge w.r.t. the query entity it is connected to,  $s$ . The initial edge representations are
$$
E _ {0} = \sigma \left(L (\mathbf {i} _ {\mathbf {r}}) \mathbf {W} _ {\mathbf {0}} ^ {(\mathbf {r})}\right) + \mathbf {E} _ {\mathrm {d i r}} [ \mathbf {i} _ {\mathrm {d i r}} ] \tag {3}
$$
where  $\mathbf{W}_0^{(\mathbf{r})}$  is a  $\mathbb{R}^{d_0\times d}$  parameter similar to  $\mathbf{W}_0^{(\mathbf{e})}$ , and  $\mathbf{E}_{\mathrm{dir}}$  is the binary embedding tables to encode the direction of the edge relative to the query entity.
4.2.3 Message Passing Layer
We follow the work of Galkin et al. (2020); Vashishth et al. (2020) and do not learn separate message passing transformations for each relation type. Instead, we learn only two transformations: One for messages passed in the forward direction and one for those passed backwards.
$$
\mathrm {M} _ {f w d} = \left(X \left[ \mathbf {i} _ {\mathbf {h}} \right] \parallel E\right) \mathbf {W} _ {\mathbf {f w d}} \tag {4}
$$
$$
\mathrm {M} _ {\text {b a c k}} = \left(X [ \mathbf {i} _ {\mathbf {t}} ] \| E\right) \mathbf {W} _ {\text {b a c k}}
$$
where  $\mathbf{W}_{<\mathrm{dir}>}$  is the  $\mathbb{R}^{d\times d}$  weight matrix for the corresponding message direction, and  $\parallel$  is the concatenation operator. The messages in  $\mathrm{M}_{fwd}$  are sent to  $\mathbf{i_t}$  while the messages in  $\mathrm{M}_{back}$  are sent to  $\mathbf{i_h}$ . We then use mean aggregation and residual connections giving an entity update of
$$
x _ {i} ^ {\prime} = \frac {1}{| \mathcal {N} (i) |} \sum_ {\mathcal {N} (i)} \mathrm {M} _ {f w d} [ \mathbf {i _ {t}} = i ] + \mathrm {M} _ {b a c k} [ \mathbf {i _ {h}} = i ]
$$
$$
x _ {i} ^ {\prime} = \sigma \left(x _ {i} ^ {\prime}\right) + x _ {i} \tag {5}
$$
This calculation can be efficiently computed in parallel using vectorized scatter_mean.
4.2.4 Edge Update Layer
Similarly to (Gong and Cheng, 2019), we also maintain and update edge representations through each layer. We do this using a simple transformation using the previous edge representation and the entity representations of the head and tail:
$$
E ^ {\prime} = \sigma \left(\left(X [ \mathbf {i} _ {\mathbf {h}} ] \| E \| X [ \mathbf {i} _ {\mathbf {t}} ]\right) \mathbf {W} ^ {(\mathbf {e})}\right) + E \tag {6}
$$
where  $\mathbf{W}^{(\mathbf{e})}$  is the  $\mathbb{R}^{3d\times d}$  edge update parameter. We use edge update layers between each message passing layer.
4.2.5 Combining the Encoders
We combine the MPNN and the language model in a sequential manner, in which the output of the language model is used to replace the features for the entities being queried (query entities and target entities). These are the same entities whose representations will eventually be used by the decoder.
4.3 Decoders
We use a simple TransE (Bordes et al., 2013) model to score each candidate triple.
For given triple  $q = (h,r,t)$ , with final entity representations  $\mathbf{X}$ , and relation embedding table  $\mathbf{H} \in \mathbb{R}^{|\mathcal{R}| \times d}$ , TransE scores the triple as
$$
T r a n s E (q) = - \| X [ h ] + \mathbf {H} [ r ] - X [ t ] \| _ {2} \tag {7}
$$
where higher scores indicate a higher likelihood of existing in the graph. While, TransE, on it's own cannot represent all classes of relations (e.g. symmetric relations), this analysis does not apply when tied to a more expressive encoder.
4.4 Scalability
We aim to be able to use StATIK in real-world, large-scale settings. As a result, scalability is critically important. We employ a number of techniques that allow us to implement our message passing neural network on very large graphs.
4.4.1 Neighbor sampling
Neighbor sampling has been a key technique in scaling GNNs to large graphs (Hamilton et al., 2017b; Chen et al., 2018; bing Huang et al., 2018; Markowitz et al., 2021) and is even more critical with knowledge graphs. In homogeneous graphs, nodes have relatively similar degrees. In knowledge graphs, this is not the case. In our largest informally tested graph some nodes have degree greater than 10 million while the vast majority have degree less than 10. As a result neighbor sampling becomes a critical step to reduce computational complexity of the model.
Relations	Train	Validation	Test
Entities	Triples	Entities	Triples	Entities	Triples
WN18RR	11	32,755	69,585	4,094	11,381	4,094	12,087
FB15k-237	237	11,633	215,082	1,454	42,164	1,454	52,870
Wikidata-5M	822	4,579,609	20,496,514	7,374	6,699	7,475	6,894
Table 2: Statistics of the datasets used in the experiments.
4.4.2 Compact Adjacency
In order to sample efficiently we adapt the compact adjacency structure from Markowitz et al. (2021) to the multi-relational setting. This data structure is similar to a CSR-formatted sparse matrix in which all the data is left-aligned. This allows fast row access, giving us the ability to get the neighbors of each node and sample from them in an efficient manner. Further details can be found in the appendix.
4.4.3 Query-Target Independence
As done in StAR and BLP, we separately calculate the embeddings for queries and targets, requiring  $\mathcal{Q} + \mathcal{N}$  (#queries + #targets) passes through the encoder to run the graph completion task. However, unlike StAR, we do not need every entity-relation pair in order to do head prediction.
4.5 Optimized Negative Sampling
Rather than sampling negative entities as targets, that require additional passes through the encoder. We use the true targets for other queries in the batch as negative samples, filtering if the target would form a true triple. This does affect inference but does dramatically accelerate training.
In this section we describe our thorough evaluation protocol. We describe the two settings in which we evaluate our inductive model and the infrastructure on which we implemented our model before demonstrating its effectiveness on three benchmark datasets. To analyze the importance of the various components of StATIK, we present the results of some detailed ablation studies.
5.1 Evaluation Protocol
Our evaluation protocol follows that in (Bordes et al., 2013), wherein for each test triple  $(h', r, t')$  we generate two queries for the tail and head prediction tasks:  $(h', r, ?)$  and  $(?, r, t')$ . For each query, we consider every entity in the graph as a candidate target. We then rank all candidate triples in decreasing order of score to see where the correct triple  $(h', r, t')$  ranks. We evaluate in the filtered setting as in Bordes et al. (2013). In the filtered setting, existing valid triples are ignored when ranking candidate targets. We report Mean Reciprocal Rank (MRR), Hits@1, Hits@3 and Hits@10 in line with literature. All our metrics are reported as the average between the head prediction and tail prediction tasks. This is more challenging than training for only one type of prediction at a time.
As in Daza et al. (2021), we demonstrate the effectiveness of StATIK in both a dynamic setting as well as a transfer setting. The dynamic setting corresponds to one in which at least one of the entities present in a given test triple  $(h', r, t')$  has not been seen by the model during training, i.e.  $(h' \notin \mathcal{E}_{train} \vee t' \notin \mathcal{E}_{train})$ . The dynamic setting represents the most likely way in which a knowledge graph grows over time. New entities that are added, connect to one or more existing entities as well as other new ones. The transfer setting represents the situation in which neither the head nor the tail already exists in the knowledge graph, i.e.  $(h' \notin \mathcal{E}_{train} \wedge t' \notin \mathcal{E}_{train})$ . This is a slightly more challenging setting and performance in this setting is a good indicator of inductivity.