SlowGuess/NAACL_2022_OCR huggingface dataset:
Jivat Neet Kaur\*  $\alpha$  , Sumit Bhatia, Milan Aggarwal, Rachit Bansal\*y, and Balaji Krishnamurthy
$^{\alpha}$  Microsoft Research, India
$^{\beta}$  Media and Data Science Research Lab, Adobe, India
$^{\gamma}$  Delhi Technological University, India
t-kaurjivat@microsoft.com, sumit.bhatia@adobe.com
milaggar@adobe.com, racbansa@gmail.com, kbalaji@adobe.com
Large transformer-based pre-trained language models have achieved impressive performance on a variety of knowledge-intensive tasks and can capture factual knowledge in their parameters. We argue that storing large amounts of knowledge in the model parameters is suboptimal given the ever-growing amounts of knowledge and resource requirements. We posit that a more efficient alternative is to provide explicit access to contextually relevant structured knowledge to the model and train it to use that knowledge. We present LM-CORE – a general framework to achieve this – that allows decoupling of the language model training from the external knowledge source and allows the latter to be updated without affecting the already trained model. Experimental results show that LM-CORE, having access to external knowledge, achieves significant and robust outperformance over state-of-the-art knowledge-enhanced language models on knowledge probing tasks; can effectively handle knowledge updates; and performs well on two downstream tasks. We also present a thorough error analysis highlighting the successes and failures of LM-CORE. Our code and model checkpoints are publicly available<sup>1</sup>.
Large pre-trained language models (PLMs) (Peters et al., 2018; Devlin et al., 2019; Brown et al., 2020) have achieved state-of-the-art performance on a variety of NLP tasks. Much of this success can be attributed to the significant semantic and syntactic information captured in the contextual representations learned by PLMs. In addition to applications requiring linguistic knowledge, PLMs have also
been useful for a variety of tasks involving factual knowledge and it has been shown that models such as BERT (Devlin et al., 2019) and T5 (Raffel et al., 2020) store significant world knowledge in their parameters (Petroni et al., 2019).
PLMs are typically fed a large amount of unstructured text which leads to the linguistic nuances and world knowledge being captured in the model parameters. This implicit storage of the knowledge in the form of the parameter weights not only leads to poor interpretability while analyzing model predictions but also poses constraints on the amount of knowledge that can be stored. It is not practical to pack all the ever-evolving world knowledge in the language model parameters due to the great financial and environmental costs incurred by training of the PLMs. Further, since the PLMs acquire knowledge from the text corpora they are trained on, they tend to become sensitive to the contextual and linguistic variations (Jiang et al., 2020). Moreover, PLMs do not contain explicit grounding to real world entities, and hence, often find it difficult to recall factual knowledge (Logan et al., 2019). For example, the model may not be able to recall correct information and successfully complete the sentence, "The birthplace of Barack Obama is ______", if the LM has seen this fact in a different context during training (e.g., "Barack Obama was born in Honolulu, Hawaii").
Large scale structured knowledge bases (KBs) such as YAGO (Suchanek et al., 2007) and Wiki-data (Vrandecic and Krötzsch, 2014) offer a rich resource of high quality structured knowledge that can provide the PLMs with explicit grounding to real world entities. Consequently, efforts have been made to integrate factual knowledge into PLMs and create entity-enhanced language models (Peters et al., 2019; Zhang et al., 2019; Sun et al., 2020; Liu et al., 2020; Wang et al., 2021a,b). However,
*Work done as an intern at the Media and Data Science Research Lab, Adobe, India.
1https://github.com/sumit-research/ lmcore
Findings of the Association for Computational Linguistics: NAACL 2022, pages 750 - 769 July 10-15, 2022 ©2022 Association for Computational Linguistics
Figure 1: Language Model Pre-Training with Contextually Relevant External Knowledge: ① Using a sentence sampled from the pre-training corpus, an input  $(x)$  is created by selecting an entity mention at random from the potential mask candidates (underlined in red). ② An NER tagger is then applied to the masked input sequence  $(x)$  to identify named entities (underlined in black). ③ For the identified entities, the Knowledge Retrieval module fetches the set  $\mathcal{T}_x$  of all the triples from the Knowledge Base and then ④ scores all the retrieved triples using input-triple and input-relation similarity (details in Section 3.2). ⑤ The top- $k$  triples are fed to the Language Model encoder along with the input sequence  $(x)$  and the model is trained to predict the masked token.
these works either update the PLM parameters or modify the architecture to facilitate the storage of factual knowledge in the model layers and parameters, making it expensive to update knowledge.
In this work, we step back and ask - what if instead of focusing on storing the knowledge in the language model parameters, we provide the model with contextually relevant external knowledge and train it to use this knowledge? This approach offers several potential advantages - (i) we can utilize the already available high-quality large-scale knowledge bases such as YAGO and Wikidata; (ii) not all the knowledge needs to be packed in the parameters of the model resulting in lighter, smaller and greener models; and (iii) as new knowledge becomes available, the knowledge base can be updated independently of the language model.
Our Contributions: We present LM-CORE, a framework for augmenting language models with contextually relevant external knowledge. The LM-CORE framework is summarized in Figure 1 and consists of a contextual knowledge retriever that fetches relevant knowledge from an external KB and passes it to the language model along with the input text. The language model is then trained with a modified entity-oriented masked language modeling objective (Section 3). Our proposed solution is simple, yet highly effective. Experiments on benchmark knowledge probes show that the pro
posed approach leads to significant performance improvements over base language models as well as state-of-the-art knowledge enhanced variants of the language models (Section 4). We find that with access to contextually relevant external knowledge, LM-CORE is less sensitive to the contextual variations in input text. We also show how LM-CORE can handle knowledge updates without any re-training and compare the performance of LM-CORE on two knowledge-intensive downstream tasks. Finally, we present an in-depth analysis of cases where our proposed approach gives incorrect answers paving the way for further research in this direction (Section 4.7).
Augmenting Additional Knowledge in PLMs: Previous works on augmenting PLMs with additional knowledge can be grouped into two categories. One line of work adopts a retrieve and read framework where the model is trained to retrieve relevant information followed by a reading comprehension step to perform the downstream task (Lee et al., 2019a; Guu et al., 2020; Agarwal et al., 2021). While our proposal has similarities with this line of work in terms of retrieving the contextual knowledge, there are two major differences. First, most of these works consider external knowledge in the form of unstructured text (such as
Wikipedia documents). However, extracting factual knowledge from unstructured text is hard and error-prone due to ambiguities in natural language and infrequent mentions of entities of interest (Peters et al., 2019). This issue can be alleviated by using a structured knowledge base where the knowledge is represented (mostly) unambiguously – each fact is a triple in the knowledge base. Further, these approaches employ explicit supervision during pretraining to train the model to fetch relevant passages from the text. This results in systems that are more complex and resource-hungry than the base PLMs used and also make it difficult to reuse or adapt the models to different sources of knowledge.
The second body of work has focused on injecting the factual knowledge directly into the model parameters by feeding more data to the model during pre-training (Poerner et al., 2020; Roberts et al., 2020). A promising direction explored recently is utilizing structured knowledge bases to augment Transformer-based LMs. ERNIE (Zhang et al., 2019) and KnowBERT (Peters et al., 2019) are notable efforts in this direction where the entity information from knowledge bases is explicitly linked with the input text during pre-training yielding entity-enhanced variants of BERT models with entity representations integrated within the Transformer layers. An alternative way of training entity-aware language models is illustrated by frameworks such as CoLAKE (Sun et al., 2020) and KEPLER (Wang et al., 2021b) that jointly learn the language and knowledge representations thereby producing language models augmented with factual knowledge and knowledge embeddings enhanced with textual context. However, these approaches, by design, will lead to larger and larger models to store the ever-growing abundant knowledge. Further, due to the strong coupling between the knowledge and language signals, updating or adding knowledge requires re-training of the model.
Examining the knowledge contained in PLMs: Petroni et al. (2019) posit that while training over large amounts of input text, PLMs may also be storing (implicit) relational knowledge in their parameters and proposed the Language Model Analysis (LAMA) framework to measure the relational knowledge stored in a PLM. Jiang et al. (2020) argue that due to the sensitivity of the PLMs on the input context, such manually created prompts are sub-optimal and might fail to retrieve facts that the PLM does know, thus providing only a lower bound
Figure 2: We create our pre-training corpus from Wikipedia by masking entity spans detected using anchor text of hyperlinks.
estimate of the knowledge contained in it. Subsequent work (Shin et al., 2020; Zhong et al., 2021) has attempted to generate better prompts in order to tighten this estimate. Poerner et al. (2020) introduced LAMA-UHN (UnHelpfulNames), a much harder subset of LAMA where the input probes provide little or no helpful contextual signals from other tokens in the probe, thus measuring the innate ability of the PLM to recall information.
3 LM-CORE: Knowledge Retrieval and Training Framework
Task setting and Overview: Consider a language model  $\mathcal{L}$  (such as BERT and RoBERTa) and a knowledge base  $\mathcal{K} = \{t_{hrt} = < h, r, t > | h, t \in \mathcal{E}; r \in \mathcal{R}\}$ . Here, we consider the knowledge base  $\mathcal{K}$  as a set of triples such that each triple  $t_{hrt}$  represents the relationship  $r$  between entities  $h$  and  $t$ .  $\mathcal{E}$  is the set of all the entities, and  $\mathcal{R}$  is the set of all the relationship types present in the knowledge base. Given a text input  $x$ , the proposed LM-CORE framework retrieves a set of triples  $\mathcal{T}_x \in \mathcal{K}$  such that the triples in  $\mathcal{T}_x$  are contextually relevant to  $x$ . The language model is then presented with the original input  $x$  and the contextually relevant knowledge in the form of  $\mathcal{T}_x$  and is trained to make predictions using this additional knowledge. We posit that the model essentially needs to learn relevant semantic associations between natural language input text and various relation types present in the knowledge base. Identifying the correct relation types will help the model leverage the corresponding relevant facts in order to make an accurate prediction. This is accomplished via a modified Masked Language Modeling (MLM) (Devlin et al., 2019) pre-training objective. Figure 1 summarizes the complete workflow of our proposed LM-CORE framework and we describe the three main components in detail in the following sub-sections.
3.1 Entity span masking
Masked Language Modelling is a popular task used for training PLMs where the objective is to predict the masked token in the input sequence. In order to improve model's grounding to real world entities, previous works have adopted different strategies for explicitly masking entity information in the input text by using entity representations obtained by knowledge base embeddings (Zhang et al., 2019), using named entity recognizers (NER) and regular expressions (Guu et al., 2020), and verbalizing knowledge base triples (Agarwal et al., 2021). These approaches often result in noisy masks due to the limitations of underlying rules, and NER and entity linking systems. To overcome these limitations, we propose a novel way of creating high-quality and accurate entity masks by using Wikipedia as the base corpus for training. Note that in order to create entity masks, we need to identify corresponding entity mentions in the input text for which we utilize the human-annotated links in Wikipedia. The official style guidelines of Wikipedia require the editors to link mentions of topics to their corresponding Wikipedia pages. In Figure 2, the left textbox shows a screenshot of the Wikipedia article about Batman where various other related topics, or concepts, are linked to their corresponding Wikipedia pages (underlined in red in the figure, and displayed as blue anchor-text in Wikipedia). This information provides us with high-quality human annotation of entity mentions in the input text. As illustrated in Figure 2, the underlined tokens (such as DC Comics, Bob Kane, Bill Finger) constitute the set of entity tokens that could be masked. For each such mask, we can also obtain the corresponding contextual knowledge from the external knowledge base (illustrated for Bob Kane in the right text box). By masking only the entity tokens (instead of randomly sampled words) and providing contextually relevant knowledge to the model retrieved from the knowledge base (as described in next subsection), we expect the model to learn to predict the masked entity tokens by utilizing the external knowledge.
3.2 Contextual Knowledge Retrieval
After preparing the masked input for training, the second component in our framework fetches contextually relevant knowledge to feed to the language model. Consider the sentence, "Warren Buffet is the chairman of [MASK]", where the masked
token is Berkshire Hathaway. In the typical MLM setting, the model only has access to the linguistic and contextual clues present in the input text to predict the masked token. However, if contextually relevant information is available as additional input, the model can use it to output the correct token.
We consider the problem of finding contextually relevant facts given the input query text as an information retrieval (IR) problem and adopt a retrieve and re-rank approach that has empirically been found to perform well in a variety of tasks (Chen et al., 2017; Wang et al., 2017; Das et al., 2019; Yang et al., 2019). Recall the example input discussed above - "Warren Buffett is the chairman of [MASK]". Intuitively, in this input text, there are two important signals that the retriever needs to utilize - entity and relation information. The entity mention Warren Buffett indicates that we need to fetch facts related to Warren Buffet from the knowledge base. Typically, there are numerous facts related to a given entity in the knowledge base, especially for popular entities such as Warren Buffett. Thus, the retriever also needs to utilize the presence of the word chairman to retrieve facts (KB triples) representing the management or executive relation.
Given an input text, our retriever pipeline performs Named Entity Recognition (NER) to identify named entity mentions in the input text. We use the NER model from FLAIR (Akbik et al., 2019) to identify named entity mentions and then select KB entities having maximum overlap with the mention-span of the identified entities. For instance, if the input query is "Buffett was born in [MASK]", all of the entities containing Buffett - Warren_Buffett, Howard_Warren_Buffett, Howard_Graham_Buffett, Volcano_(Jimmy_Buffett-song) etc. are selected, but if the query is "Warren Buffett was born in [MASK]", only the first two entities will be chosen). Once these entities are selected, all the facts from the KB involving these entities are retrieved (denoted by  $\mathcal{T}_x$  in Figure 1).
After retrieving the facts involving the entities mentioned in the input, we next need to rank these triples based on their relevance to the input. In order to measure the contextual relevance of a given triple  $t$  to the input  $x$ , we compute the following two scores.
Query-Triple similarity: We obtain representations of the input text  $x$  as well as the triple  $t$  and compute the inner product of the representations to
obtain the similarity score as follows.
$$
\operatorname {s i m} (x, t) = \operatorname {E m b} (x) ^ {T} \operatorname {E m b} (t), t \in \mathcal {T} _ {x} \tag {1}
$$
Here,  $\mathsf{Emb}(\cdot)$  is obtained using the Sentence Transformer (Reimers and Gurevych, 2019). While it is straightforward to obtain representations of input  $x$ , sentence transformer can not be applied directly to KB triples. Application of KB embeddings such as TransE (Bordes et al., 2013) is also not feasible as then the representations of the input text and triples will be in different embedding spaces. To overcome this, we adopt a simple approach of verbalizing the knowledge base triples by concatenating the head entity, relationship and the tail entity, and obtain the representation of the verbalized triple from the sentence transformer. For example, the triple (Warren_Buffett, hasOccupation, Investor) is verbalized as Warren_Buffett has occupation Investor and is fed as input to the sentence transformer.
Relation-based scoring: A triple is highly relevant for the input text if the triple represents the same relationship that is being talked about in the text. To capture this intuition, we embed all the relation types in the KB in the same embedding space as triples using the sentence transformer and compute the similarity between the input text and the relation type of the triple as follows.
$$
s i m (x, r) = \mathsf {E m b} (x) ^ {T} \mathsf {E m b} (r), r \in \mathcal {R} \quad (2)
$$
where  $\mathcal{R}$  is the set of all relations in the KB. The final relevance score for the triple  $t$ , relevance  $(x, t)$  is obtained by taking a product of the above two scores. Based on this final score, we select the top- $k$  triples that constitute the contextual knowledge to be fed as input along with  $x$  to the LM. We use  $k = 8$  in this work (See Appendix 4.3 for effect of varying  $k$ ). Some illustrative examples of the final retrieved knowledge base triples are presented in Appendix A.3.
3.3 Language Model Pre-training with Contextual Knowledge
With the masked training corpus and the module to fetch contextually relevant knowledge, we now train the model to utilize the additional contextual knowledge to predict the masked token. From the masked corpus, we select a sentence and a valid entity span is chosen at random out of all the potential spans in the sentence. We mask this span to create the input text  $x$ . We filter out sentences starting with pronouns such as he, she, her, and
they as we observed that most of such sentences do not contain other useful signals to unambiguously predict the masked words. For instance, if the input example is - "He developed an interest in investing in his youth, eventually entering the Wharton School of the University of Pennsylvania" and Wharton School of the University of Pennsylvania is masked, the remaining words in the sentence are not providing any informative signals to the model to predict the masked tokens. Given the input sentence thus selected, the contextual knowledge retriever fetches the relevant triples from the knowledge base. The representations of the input sentence and the retrieved triples are then concatenated and fed to the model and the model is trained to minimize the following MLM loss.
$$
L _ {M L M} = \frac {1}{M} \sum_ {m = 1} ^ {M} \log p \left(x _ {i n d _ {m}} \mid x, t _ {1}, t _ {2}, \dots , t _ {k}\right) \tag {3}
$$
where  $M$  is the total number of [MASK] tokens in  $x$  and  $ind_{m}$  is the index of the  $m^{th}$  masked token.
With the additional contextual information available to the model, we expect the model to learn the associations between linguistic cues in the input text and relevant relationship information in the triples. For example, we expect the model to associate different ways in which someone's date of birth could be mentioned in natural language (such as  $X$  was born on, the birthday of  $X$  is, and numerous other linguistic variations) to the KB relation birthDate and utilize the information from the corresponding triple. Note that since the types of relations in the knowledge base are relatively small in number, and do not change often, we expect the model to generalize well and be more robust to linguistic variations.
4 Experiments and Discussions
Data Sources and Pre-processing: We create our pre-training corpus using the December 20, 2018 snapshot of English Wikipedia that contains about  $5.5\mathrm{M}$  documents. Processing all the articles following the masking strategy described in Section 3.1 resulted in a total of  $\sim 46.3\mathrm{M}$  sentences with valid masks, from which we randomly sample sentences to create input examples.
In order to illustrate the general nature of LM-CORE, we used two different PLMs as our LM encoders - BERT-base (uncased) model and RoBERTa-base (cased) model. We use YAGO and
Table 1: Mean precision at one  $(\mathrm{P}@\mathrm{l})$  of various models on LAMA probe. We group all the models based on the base language model used (BERT or RoBERTa). For LM-CORE,  $(\cdot ,\cdot)$  indicates the variant - (b, r corresponds to BERT and RoBERTa, respectively, and y, w indicate YAGO and Wikidata5M, respectively). Best results in each column are highlighted in bold and the second best performance is underlined
Complete	Google-RE	T-REx	SQuAD	Concept Net
DoB	PoB	PoD	All	1-1	N-1	N-M	All
BERT-based models
BERT-base	24.73	1.59	15.46	10.33	9.12	67.94	32.67	23.54	30.83	14.29	15.88
BERT-large	25.44	1.59	15.53	12.16	9.76	74.23	31.30	25.30	31.05	17.61	18.72
ERNIE	22.16	1.42	13.48	4.97	6.62	61.51	28.57	21.93	27.58	13.62	14.83
LM-CORE(b,y)	39.64	64.44	52.71	50.98	56.04	74.37	51.18	34.57	45.83	15.61	14.78
LM-CORE(b,w)	42.83	0.66	37.62	31.11	23.13	81.79	59.86	45.48	55.32	17.28	16.15
RoBERTa-based models
RoBERTa	20.46	1.85	12.98	1.23	5.35	57.49	23.14	21.59	24.21	12.94	18.47
RoBERTa-large	24.24	1.41	12.48	0.46	4.78	70.24	29.08	23.28	28.82	18.88	22.09
KEPLER	19.36	1.47	11.73	3.08	5.43	52.32	21.58	21.41	23.01	9.10	17.25
CoLAKE	23.38	1.79	15.72	10.79	9.43	64.08	29.40	23.54	28.80	8.39	17.17
LM-CORE(r,y)	34.60	46.33	43.47	26.35	38.71	68.21	45.30	30.40	40.60	13.29	17.53
LM-CORE(r,w)	41.96	0.38	33.11	28.20	20.56	70.21	60.30	43.18	54.11	15.73	18.38
Wikidata as two different knowledge bases giving us four variants of LM-CORE ( $\{bert, roberta\} \times \{yago, wikidata\}$ ). We use the English Wikipedia version of YAGO 4 (Suchanek et al., 2007) and preprocess it to obtain our retrieval corpus consisting of roughly 17M triples spanning over 4.9M entities and 131 unique relations. For Wikidata, we used the Wikidata5M version (Wang et al., 2021b) that consists of roughly 21M triples covering 821 unique relations and 4.8M entities. Further details regarding retrieval corpus generation and processing can be found in the Appendix (Section A.2). For computing triple representations for retrieval (Section 3.2), we concatenate the subject (head), relation, and object (tail) of triples and embed them using the Sentence Transformers (Reimers and Gurevych, 2019) and obtain the 768-dimensional embeddings (same as LM encoder dimensions).
4.1 Does External Knowledge Help PLMs in Knowledge Intensive Tasks?