How have researchers attempted to enable large language models to make use of the information stored in knowledge graphs?
Local context approaches focus on entity men-
tions and external knowledge in individual sen-
tences to enable fine-grained knowledge inclusion.
A straightforward way is to encode KG entities
with KG embeddings (Bordes et al., 2013; Lin et al.,
2015; Cucala et al., 2021; Sun et al., 2018) and in-
fuse the embeddings with language representations
(Hu et al., 2021; Feng et al., 2021; Kang et al.,
2022). Later approaches focus on augmenting pre-
trained LMs with KGs by introducing knowledge-
aware training tasks and LM architectures (Wang
et al., 2021b,a; Sridhar and Yang, 2022; Moiseev
et al., 2022; Kaur et al., 2022; Hu et al., 2022; Arora
et al., 2022; de Jong et al., 2021; Meng et al., 2021;
He et al., 2021). Topic models were also introduced
to enrich document representation learning (Gupta
et al., 2018; Chaudhary et al., 2020; Wang et al.,
2018). However, local context approaches fall short
of leveraging inter-sentence and inter-entity knowl-
edge, resulting in models that could not grasp the
full picture of the text-knowledge interactions
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
Infusing factual knowledge into pre-trained models is fundamental for many knowledge-intensive tasks. In this paper, we proposed Mixture-of-Partitions (MoP), an infusion approach that can handle a very large knowledge graph (KG) by partitioning it into smaller sub-graphs and infusing their specific knowledge into various BERT models using lightweight adapters. To leverage the overall factual knowledge for a target task, these sub-graph adapters are further fine-tuned along with the underlying BERT through a mixture layer. We evaluate our MoP with three biomedical BERTs (SciBERT, BioBERT, PubmedBERT) on six downstream tasks (inc. NLI, QA, Classification), and the results show that our MoP consistently enhances the underlying BERTs in task performance, and achieves new SOTA performances on five evaluated datasets.
[{"first":"Zaiqiao","last":"Meng","middle":[]},{"first":"Fangyu","last":"Liu","middle":[]},{"first":"Thomas","last":"Clark","middle":[]},{"first":"Ehsan","last":"Shareghi","middle":[]},{"first":"Nigel","last":"Collier","middle":[]}]
LM-CORE: Language Models with Contextually Relevant External Knowledge
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 sub-optimal 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 \textit{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.
[{"first":"Jivat","last":"Kaur","middle":["Neet"]},{"first":"Sumit","last":"Bhatia","middle":[]},{"first":"Milan","last":"Aggarwal","middle":[]},{"first":"Rachit","last":"Bansal","middle":[]},{"first":"Balaji","last":"Krishnamurthy","middle":[]}]
SKILL: Structured Knowledge Infusion for Large Language Models
Large language models (LLMs) have demonstrated human-level performance on a vast spectrum of natural language tasks. However, it is largely unexplored whether they can better internalize knowledge from a structured data, such as a knowledge graph, or from text. In this work, we propose a method to infuse structured knowledge into LLMs, by directly training T5 models on factual triples of knowledge graphs (KGs). We show that models pre-trained on Wikidata KG with our method outperform the T5 baselines on FreebaseQA and WikiHop, as well as the Wikidata-answerable subset of TriviaQA and NaturalQuestions. The models pre-trained on factual triples compare competitively with the ones on natural language sentences that contain the same knowledge. Trained on a smaller size KG, WikiMovies, we saw 3x improvement of exact match score on MetaQA task. The proposed method has an advantage that no alignment between the knowledge graph and text corpus is required in curating training data. This makes our method particularly useful when working with industry-scale knowledge graphs.
[{"first":"Fedor","last":"Moiseev","middle":[]},{"first":"Zhe","last":"Dong","middle":[]},{"first":"Enrique","last":"Alfonseca","middle":[]},{"first":"Martin","last":"Jaggi","middle":[]}]
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Tuning pre-trained language models (PLMs) with task-specific prompts has been a promising approach for text classification. Particularly, previous studies suggest that prompt-tuning has remarkable superiority in the low-data scenario over the generic fine-tuning methods with extra classifiers. The core idea of prompt-tuning is to insert text pieces, i.e., template, to the input and transform a classification problem into a masked language modeling problem, where a crucial step is to construct a projection, i.e., verbalizer, between a label space and a label word space. A verbalizer is usually handcrafted or searched by gradient descent, which may lack coverage and bring considerable bias and high variances to the results. In this work, we focus on incorporating external knowledge into the verbalizer, forming a knowledgeable prompttuning (KPT), to improve and stabilize prompttuning. Specifically, we expand the label word space of the verbalizer using external knowledge bases (KBs) and refine the expanded label word space with the PLM itself before predicting with the expanded label word space. Extensive experiments on zero and few-shot text classification tasks demonstrate the effectiveness of knowledgeable prompt-tuning.
[{"first":"Shengding","last":"Hu","middle":[]},{"first":"Ning","last":"Ding","middle":[]},{"first":"Huadong","last":"Wang","middle":[]},{"first":"Zhiyuan","last":"Liu","middle":[]},{"first":"Jingang","last":"Wang","middle":[]},{"first":"Juanzi","last":"Li","middle":[]},{"first":"Wei","last":"Wu","middle":[]},{"first":"Maosong","last":"Sun","middle":[]}]
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Natural language understanding tasks such as open-domain question answering often require retrieving and assimilating factual information from multiple sources. We propose to address this problem by integrating a semi-parametric representation of a large text corpus into a Transformer model as a source of factual knowledge. Specifically, our method represents knowledge with `mention memory', a table of dense vector representations of every entity mention in a corpus. The proposed model - TOME - is a Transformer that accesses the information through internal memory layers in which each entity mention in the input passage attends to the mention memory. This approach enables synthesis of and reasoning over many disparate sources of information within a single Transformer model. In experiments using a memory of 150 million Wikipedia mentions, TOME achieves strong performance on several open-domain knowledge-intensive tasks, including the claim verification benchmarks HoVer and FEVER and several entity-based QA benchmarks. We also show that the model learns to attend to informative mentions without any direct supervision. Finally we demonstrate that the model can generalize to new unseen entities by updating the memory without retraining.
[{"first":"Michiel","last":"Jong","middle":["de"]},{"first":"Yury","last":"Zemlyanskiy","middle":[]},{"first":"Nicholas","last":"FitzGerald","middle":[]},{"first":"Fei","last":"Sha","middle":[]},{"first":"William","last":"Cohen","middle":[]}]
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
Pre-trained language representation models (PLMs) learn effective language representations from large-scale unlabeled corpora. Knowledge embedding (KE) algorithms encode the entities and relations in knowledge graphs into informative embeddings to do knowledge graph completion and provide external knowledge for various NLP applications. In this paper, we propose a unified model for Knowledge Embedding and Pre-trained LanguagE Representation (KEPLER), which not only better integrates factual knowledge into PLMs but also effectively learns knowledge graph embeddings. Our KEPLER utilizes a PLM to encode textual descriptions of entities as their entity embeddings, and then jointly learn the knowledge embeddings and language representations. Experimental results on various NLP tasks such as the relation extraction and the entity typing show that our KEPLER can achieve comparable results to the state-of-the-art knowledge-enhanced PLMs without any additional inference overhead. Furthermore, we construct Wikidata5m, a new large-scale knowledge graph dataset with aligned text descriptions, to evaluate KE embedding methods in both the traditional transductive setting and the challenging inductive setting, which needs the models to predict entity embeddings for unseen entities. Experiments demonstrate our KEPLER can achieve good results in both settings.
[{"first":"Xiaozhi","last":"Wang","middle":[]},{"first":"Tianyu","last":"Gao","middle":[]},{"first":"Zhaocheng","last":"Zhu","middle":[]},{"first":"Zhiyuan","last":"Liu","middle":[]},{"first":"Juanzi","last":"Li","middle":[]},{"first":"Jian","last":"Tang","middle":[]}]
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Popular language models (LMs) struggle to capture knowledge about rare tail facts and entities. Since widely used systems such as search and personal-assistants must support the long tail of entities that users ask about, there has been significant effort towards enhancing these base LMs with factual knowledge. We observe proposed methods typically start with a base LM and data that has been annotated with entity metadata, then change the model, by modifying the architecture or introducing auxiliary loss terms to better capture entity knowledge. In this work, we question this typical process and ask to what extent can we match the quality of model modifications, with a simple alternative: using a base LM and only changing the data. We propose metadata shaping, a method which inserts substrings corresponding to the readily available entity metadata, e.g. types and descriptions, into examples at train and inference time based on mutual information. Despite its simplicity, metadata shaping is quite effective. On standard evaluation benchmarks for knowledge-enhanced LMs, the method exceeds the base-LM baseline by an average of 4.3 F1 points and achieves state-of-the-art results. We further show the gains are on average 4.4x larger for the slice of examples containing tail vs. popular entities.
[{"first":"Simran","last":"Arora","middle":[]},{"first":"Sen","last":"Wu","middle":[]},{"first":"Enci","last":"Liu","middle":[]},{"first":"Christopher","last":"Re","middle":[]}]
Explaining Toxic Text via Knowledge Enhanced Text Generation
Warning: This paper contains content that is offensive and may be upsetting.Biased or toxic speech can be harmful to various demographic groups. Therefore, it is not only important for models to detect these speech, but to also output explanations of why a given text is toxic. Previous literature has mostly focused on classifying and detecting toxic speech, and existing efforts on explaining stereotypes in toxic speech mainly use standard text generation approaches, resulting in generic and repetitive explanations. Building on these prior works, we introduce a novel knowledge-informed encoder-decoder framework to utilize multiple knowledge sources to generate implications of biased text.Experiments show that our knowledge informed models outperform prior state-of-the-art models significantly, and can generate detailed explanations of stereotypes in toxic speech compared to baselines, both quantitatively and qualitatively.
[{"first":"Rohit","last":"Sridhar","middle":[]},{"first":"Diyi","last":"Yang","middle":[]}]
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
Interactions between entities in knowledge graph (KG) provide rich knowledge for language representation learning. However, existing knowledge-enhanced pretrained language models (PLMs) only focus on entity information and ignore the ﬁne-grained relationships between entities. In this work, we propose to incorporate KG (including both entities and relations) into the language learning process to obtain KG-enhanced pretrained Language Model, namely KLMo. Speciﬁcally, a novel knowledge aggregator is designed to explicitly model the interaction between entity spans in text and all entities and relations in a contextual KG. An relation prediction objective is utilized to incorporate relation information by distant supervision. An entity linking objec-tive is further utilized to link entity spans in text to entities in KG. In this way, the structured knowledge can be effectively integrated into language representations. Experimental results demonstrate that KLMo achieves great improvements on several knowledge-driven tasks, such as entity typing and relation clas-siﬁcation, comparing with the state-of-the-art knowledge-enhanced PLMs.
[{"first":"Lei","last":"He","middle":[]},{"first":"Suncong","last":"Zheng","middle":[]},{"first":"Tao","last":"Yang","middle":[]},{"first":"Feng","last":"Zhang","middle":[]}]
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
Leveraging factual knowledge to augment pretrained language models is of paramount importance for knowledge-intensive tasks, such as question answering and fact checking (Petroni et al., 2021). Especially in the biomedical domain where public training corpora are limited and noisy, trusted biomedical KGs are crucial for deriving accurate inferences (Li et al., 2020;. However, the infusion of knowledge from realworld biomedical KGs, where entity sets are very large (e.g. UMLS, Bodenreider 2004, contains ∼4M entities) demands highly scalable solutions.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
Although many general knowledge-enhanced language models have been proposed, most of them rely on a computationally expensive joint training of an underlying masked language model (MLM) along with a knowledge-infusion objective function to minimize the risk of catastrophic forgetting (Xiong et al., 2019;Wang et al., 2021Peters et al., 2019;Yuan et al., 2021). Alternatively, entity masking (or entity prediction) has emerged as one of the most popular self-supervised training objectives for infusing entity-level knowledge into pretrained models Yu et al., 2020;He et al., 2020). However, due to the large number of entities in biomedical KGs, computing an exact softmax over all entities is very expensive for training and predicting (De Cao et al., 2021). Although negative sampling techniques could alleviate the computational issue (Sun et al., 2020), tuning an appropriately hard set of negative instances can be challenging and predicting a very large number of labels may generalize poorly (Hinton et al., 2015).
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
To address the aforementioned challenges, we propose a novel knowledge infusion approach, named Mixture-of-Partitions (MoP), to infuse factual knowledge based on partitioned KGs into pretrained models (BioBERT, Lee et al. 2020;SciB-ERT, Beltagy et al. 2019;and PubMedBERT, Gu et al. 2020). More concretely, we first partition a KG into several sub-graphs each containing a disjoint subset of its entities by using the METIS algorithm (Karypis and Kumar, 1998), and then the Transformer ADAPTER module (Houlsby et al., 2019;Pfeiffer et al., 2020b) is applied to learn portable knowledge parameters from each subgraph. In particular, using ADAPTER module to infuse knowledge does not require fine-tuning the parameters of the underlying BERTs, which is more flexible and efficient while avoiding the catastrophic forgetting issue. To utilise the independently learned knowledge from sub-graph adapters, we introduce mixture layers to automatically route useful knowledge from these adapters to downstream tasks. Figure 1 illustrates our approach.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
Our results and analyses indicate that our "divide and conquer" partitioning strategy effectively preserves the rich information presented in two biomedical KGs from UMLS while enabling us to ⇥ 0 < l a t e x i t s h a 1 _ b a s e 6 4 = " m M y N q P N R l i j
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
x / A H z u c P m 2 2 P r Q = = < / l a t e x i t > G < l a t e x i t s h a 1 _ b a s e 6 4 = " 4 e T D x P S F n r l 0 G N L y q 4 0 W Z F H J s q k = " > A A A B 8 n i c b V D L S g M x F L 1 T X 7 W + q i 7 d B I v g q s x I R R c u C i 5 0 W c E + Y D q U T J p p Q z P J k G S E M v Q z 3 L h Q x K 1 f 4 8 6 / M d P O Q l s P B A 7 n 3 E v O P W H C m T a u + + 2 U 1 t Y 3 N r f K 2 5 W d 3 b 3 9 g + r h U U f L V B H a J p J L 1 Q u x p p w J 2 j b M c N p L F M V x y G k 3 n N z m f v e J K s 2 k e D T T h A Y x H g k W M Y K N l f x + j M 2 Y Y J 7 d z Q b V m l t 3 5 0 C r x C t I D Q q 0 B t W v / l C S N K b C E I 6 1 9 j 0 3 M U G G l W G E 0 1 m l n 2 q a Y D L B I + p b K n B M d Z D N I 8 / Q m V W G K J L K P m H Q X P 2 9 k e F Y 6 2 k c 2 s k 8 o l 7 2 c v E / z 0 9 N d B 1 k T C S p o Y I s P o p S j o x E + f 1 o y B Q l h k 8 t w U Q x m x W R M V a Y G N t S x Z b g L Z + 8 S j o X d a 9 R v 3 x o 1 J o 3 R R 1 l O I F T O A c P r q A J 9 9 C C N h C Q 8 A y v 8 O Y Y 5 8 V 5 d z 4 W o y W n 2 D m G P 3 A + f w B 6 N p F f < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " W E p S v Z 3 x b z / K W 4 o T D D o c + S b x r b E = " > A A A B 9 H i c b V C 7 S g N B F L 3 r M 8 Z X 1 N J m M A h W Y V d E L S w C F l p G M A 9 I l j A 7 u U m G z M 6 u M 7 O B s O Q 7 b C w U s f V j 7 P w b Z 5 M t N P H A w O G c e 7 l n T h A L r o 3 r f j s r q 2 v r G 5 u F r e L 2 z u 7 e f u n g s K G j R D G s s 0 h E q h V Q j Y J L r B t u B L Z i h T Q M B D a D 0 W 3 m N 8 e o N I / k o 5 n E 6 I d 0 I H m f M 2 q s 5 H d C a o a M i v R u 2 h 1 1 S 2 W 3 4 s 5 A l o m X k z L k q H V L X 5 1 e x J I Q p W G C a t 3 2 3 N j 4 K V W G M 4 H T Y i f R G F M 2 o g N s W y p p i N p P Z 6 G n 5 N Q q P d K P l H 3 S k J n 6 e y O l o d a T M L C T W U i 9 6 G X i f 1 4 7 M f 1 r P + U y T g x K N j / U T w Q x E c k a I D 2 u k B k x s Y Q y x W 1 W w o Z U U W Z s T 0 V b g r f 4 5 W X S O K 9 4 l x X v 4 a J c v c n r K M A x n M A Z e H A F V b i H G t S B w R M 8 w y u 8 O W P n x X l 3 P u a j K 0 6 + c w R / 4 H z + A P 8 I k j s = < / l a t e x i t > G k < l a t e x i t s h a 1 _ b a s e 6 4 = " 2 V O S V m b W C 0 n I t n d i l P D V p U Q s O D s = " > A A A B / X i c b V D L S s N A F L 2 p r 1 p f 8 b F z E y y C q 5 K I q A s X B R e 6 r G A f 0 I Q w m U 7 a o Z N J m J k I N Q R / x Y 0 L R d z 6 H + 7 8 G y d t F t p 6 Y O B w z r 3 c M y d I G J X K t r + N y t L y y u p a d b 2 2 s b m 1 v W P u 7 n V k n A p M 2 j h m s e g F S B J G O W k r q h j p J Y K g K G C k G 4 y v C 7 / 7 Q I S k M b 9 X k 4 R 4 E R p y G l K M l J Z 8 8 8 B t j a i f u R F S I 4 x Y d p P 7 4 9 w 3 6 3 b D n s J a J E 5 J 6 l C i 5 Z t f 7 i D G a U S 4 w g x J 2 X f s R H k Z E o p i R v K a m 0 q S I D x G Q 9 L X l K O I S C + b p s + t Y 6 0 M r D A W + n F l T d X f G x m K p J x E g Z 4 s U s p 5 r x D / 8 / q p C i + 9 j P I k V Y T j 2 a E w Z Z a K r a I K a 0 A F w Y p N N E F Y U J 3 V w i M k E F a 6 s J o u w Z n / 8 i L p n D a c 8 4 Z z d 1 Z v X p V 1 V O E Q j u A E H L i A J t x C C 9 q A 4 R G e 4 R X e j C f j x X g 3 P m a j F a P c 2 Y c / M D 5 / A O + G l Y Y = < / l a t e x i t > Gk < l a t e x i t s h a 1 _ b a s e 6 4 = " b T 0 7 z 5 V + A L q Y / Q M 7 N 3 a g + J k c z v Q = " > A A A B / X i c b V D L S s N A F L 3 x W e s r P n Z u B o v g q i R F 1 I W L g g t d V r A P a E K Y T C f t 0 M m D m Y l Q Q / B X 3 L h Q x K 3 / 4 c 6 / c d J m o a 0 H B g 7 n 3 M s 9 c / y E M 6 k s 6 9 t Y W l 5 Z X V u v b F Q 3 t 7 Z 3 d s 2 9 / Y 6 M U 0 F o m 8 Q 8 F j 0 f S 8 p Z R N u K K U 5 7 i a A 4 9 D n t + u P r w u 8 + U C F Z H N 2 r S U L d E A 8 j F j C C l Z Y 8 8 9 B p j Z i X O S F W I 4 J 5 d p N 7 j d w z a 1 b d m g I t E r s k N S j R 8 s w v Z x C T N K S R I h x L 2 b e t R L k Z F o o R T v O q k 0 q a Y D L G Q 9 r X N M I h l W 4 2 T Z + j E 6 0 M U B A L / S K F p u r v j Q y H U k 5 C X 0 8 W K e W 8 V 4 j / e f 1 U B Z d u x q I k V T Q i s 0 N B y p G K U V E F G j B B i e I T T T A R T G d F Z I Q F J k o X V t U l 2 P N f X i S d R t 0 + r 9 t 3 Z 7 X m V V l H B Y 7 g G E 7 B h g t o w i 2 0 o A 0 E H u E Z X u H N e D J e j H f j Y z a 6 Z J Q 7 B / A H x u c P m O m V T Q = = < / l a t e x i t >
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
… < l a t e x i t s h a 1 _ b a s e 6 4 = " N X O y a F v B w k D t 7 o R D 3 f 9 t O R n v H b k = " > A A A B / X i c b V D L S s N A F L 2 p r 1 p f 8 b F z E y y C q 5 K I q A s X B R e 6 r G A f 0 I Q w m U 7 a o Z N J m J k I N Q R / x Y 0 L R d z 6 H + 7 8 G y d t F t p 6 Y O B w z r 3 c M y d I G J X K t r + N y t L y y u p a d b 2 2 s b m 1 v W P u 7 n V k n A p M 2 j h m s e g F S B J G O W k r q h j p J Y K g K G C k G 4 y v C 7 / 7 Q I S k M b 9 X k 4 R 4 E R p y G l K M l J Z 8 8
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
< l a t e x i t s h a 1 _ b a s e 6 4 = " L a u w r u 6 B F q z C y Q f H u g 7 S z f v z j l s = "
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
k 2 {1, 2, · · · , K} Repeat for all K sub-graphs scale up training on these very large graphs. Additionally, we observe that while individual adapters specialize towards sub-graph specific knowledge, MoP can effectively utilise their individual expertise to enhance the performance of our tested biomedical BERTs on six downstream tasks, where five of them achieve new SOTA performances.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
We denote a KG as a collection of ordered triples
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
where E and R are the sets of entities and relations, respectively. All the entities and relations are associated with their textual surface forms, which can be a single word (e.g. fever), a compound (e.g. sars-cov-2), or a short phrase (e.g. has finding site). Given a pretrained model Θ 0 , our task is to learn Φ G based on an input knowledge graph G, such that it encapsulates the knowledge from G. The training objective, L G , can be implemented in many ways such as relation classification , entity linking (Peters et al., 2019), next sentence prediction (Goodwin and Demner-Fushman, 2020), or entity prediction . In this paper, we focus on entity prediction, one of the most widely used objectives, and leave exploration of other objectives for future work.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
As mentioned earlier, exact softmax over all entities is extremely expensive (Mikolov et al., 2013;De Cao et al., 2021) for large-scale KGs, hence we resort to the principle of "divide and conquer", and propose a novel approach called Mixture-of-Partition (MoP). Specifically, our MoP first partitions a large KG into smaller sub-graphs (i.e., G → {G 1 , G 2 , . . . , G K }, §2.1), and learns subgraph specific parameters on each sub-graph separately (i.e., {Φ G 1 , Φ G 2 , . . . , Φ G K }, §2.2). Then these sub-graph parameters are fine-tuned through mixture layers to route the sub-graph specific knowledge into a target task ( §2.3).
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
Graph partitioning (i.e., partitioning the node set into mutually exclusive groups) is a critical step to our approach, since we need to properly and automatically cluster knowledge triples for supporting data parallelism and controlling computation. In particular, it must satisfy the following goals: (1) maximize the number of resulting knowledge triples to retain as much factual knowledge as possible;
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
(2) balance nodes over partitions to reduce the overall parameters across different entity prediction heads; (3) efficiency at scale for handling large KGs. In fact, an exact solution to (1) and (2) is referred to as the balanced graph partition problem, which is NP-complete. We use the METIS (Karypis and Kumar, 1998) algorithm as an approximation, simultaneously meeting all the above three requirements. METIS can handle billion-scale graphs by successively coarsening a large graph into smaller graphs, processing them quickly and then projecting the partitions back onto the larger graph, and has been used in many tasks (Chiang et al., 2019;Defferrard et al., 2016;.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
Once the large knowledge graph is partitioned, we use ADAPTER modules to infuse the factual knowledge into a pretrained Transformer model by training an entity prediction objective for each sub-graph. ADAPTERs (Houlsby et al., 2019;Pfeiffer et al., 2020b) are newly initialized modules inserted between the Transformer layers of a pretrained model. The training of ADAPTER does not require fine-tuning the existing parameters of the pretrained model. Instead, only the parameters within the ADAPTER modules are updated. In  (Gu et al., 2020) 60.24 (Gu et al., 2020) 87.56 (Gu et al., 2020) 90.32 (Nentidis et al., 2020) 36.70 (Jin et al., 2020) 83.80 (Peng et al., 2019)  this paper, we use the ADAPTER module configured by Pfeiffer et al. (2020a), which is shown in Figure 1 (b). In particular, given a sub-graph G k , we remove the tail entity name for each triple (h, r, t) ∈ G k , and transform the triple into a list of tokens:
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
The sub-graph specific ADAPTER module is trained to predict the tail entity using the representation of the [CLS] token and the parameters Φ G k are optimized by minimizing the cross-entropy loss. During the finetuning of downstream tasks, both the parameters of ADAPTER and pre-trained LM will be updated.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
Given a set of knowledge-encapsulated adapters, we use AdapterFusion mixture layers to combine knowledge from different adapters for downstream tasks. AdapterFusion is a recently proposed model (Pfeiffer et al., 2020a) that learns to combine the information from a set of task adapters by a softmax attention layer. It learns a contextual mixture weight over adapters at layer l using an attention with the softmax weights:
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
where s l,k is used to mix the adapter outputs to be passed into the next layer, and the final layer L is used to predict a task label y:
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
where f is the target task prediction head. Closely related to ours is the sparsely-gated Mixture-of-Experts layer (Shazeer et al., 2017). Alternatively, a more flexible mechanism such as Gumbel-Softmax (Jang et al., 2017) can be used for obtaining more discrete/continuous mixture weights. However, we found both alternatives underperform AdapterFusion (see Appendix for a comparison).  We evaluate our proposed MoP on two KGs, named SFull and S20Rel, which are extracted from the large biomedical knowledge graph UMLS (Bodenreider, 2004) under the SNOMED CT, US Edition vocabulary. The SFull KG contains the full relations and entities of SNOMED 2 , while the S20Rel KG is a sub-set of SFull that only contains the top 20 most frequent relations. Note that since some relations in SFull are the reversed mappings of the same entity pairs, e.g. "A has causative agent B" and "B causative agent of A", therefore for S20Rel we exclude those reversed relations in the top 20 relations. Table 2 shows the statistics of the two KGs and the used 20 relations of the S20Rel are listed in the appendix.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
We evaluate our MoP on six datasets over various downstream tasks, including four question answering (i. a binary/multiclass classification tasks. See Appendix for the detailed description of these tasks and their datasets.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
We experiment with three biomedical pretrained models, namely BioBERT (Lee et al., 2020), SciB-ERT (Beltagy et al., 2019) and PubMedBERT (Gu et al., 2020), as our base models, which have shown strong progress in biomedical text mining tasks. We first partition our KGs into different number of sub-graphs (i.e. {5, 10, 20, 40}), then for each sub-graph, we train the base models loaded with the newly initialized ADAPTER modules (with a compression rate CRate = 8) for 1-2 epochs by minimizing the cross-entropy loss. AdamW (Loshchilov and Hutter, 2018) is used as our training optimizer, and the learning rates for all the sub-graphs are fixed to 1e − 4, as suggested by (Pfeiffer et al., 2020b). Unless specified otherwise, all the reported performances are based on a partition of 20 sub-graphs, since this was optimal for task performance (see Section 3.7 for the performances over different number of partitions.).
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
In Figure 2 we report the average performance (10 runs) of the knowledge-infused PubMedBERT on two QA datasets over partitioned SFull. We can see that partitions contribute to various degrees while some (e.g. #5) have a negligible benefit. However,   Tasks   Table 1 shows the overall performance of our MoP deployed on SciBERT, BioBERT and Pub-MedBERT pretrained models. We see that MoP pretrained on the SFull KG improves both the BioBERT and PubMedBERT models for all the tasks, while the SciBERT model can be also improved on 4 out of 6 tasks. The result shows that MoP pretrained with the S20Rel KG achieves new SOTA performances on four tasks. This suggests further pruning of the knowledge triples helps task performance by reducing noise, and is a promising direction to explore in future.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
We design a controlled random partitioning scheme to test whether METIS can produce high quality partitions for training. We fix the entity size for a 20-partitioned result produced by METIS, and randomly shuffle a percentage (ranging from 0%-100%) of entities across all the sub-graphs. Table  3 shows the number of training triples numbers over different shuffling ratios. In Figure 3 we report the results on BioASQ7b and PubMedQA under different shuffling rates. We can see that the performances of MoP on both datasets degrades significantly as the shuffling rate increases, which highlights the quality of the produced partitions.     Table 4 shows the performance of PubMed-BERT+MoP trained on the SFull knowledge graph over different number of partitions. We can clearly see that under 20 partitions, PubMedBERT+MoP performs the best in both of the BioASQ7b and PubMedQA datasets, and an average entity size of 15k-30k for the sub-graphs usually yields better performance than others.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
In Figure 4, we show six examples (contexts are omitted for brevity) from BioASQ8b and compare their mixture weights of the final layer [CLS] token inferred by the PubMedBERT+MoP (SFull) model. We see that each question elicits different mixture weights, indicating that MoP can leverage the expertise of different sub-graphs depending on the target example. We also plot the word cloud over six groups of sub-graphs that are clustered by k-means according to the entity name's TF-IDF  feature of these sub-graphs. We can observe that MoP identifies the most related sub-graphs for each example (e.g. Q2 has more weights on sub-graphs [1,13,14], which specialise in 'tumor' knowledge). This validates the effectiveness of our MoP in balancing useful knowledge across adapters.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
In this paper, we proposed MoP, a novel approach for infusing knowledge by partitioning knowledge graphs into smaller sub-graphs. We show that while the knowledge-encapsulated adapters perform very differently over different sub-graphs, our proposed MoP can automatically leverage and balance the useful knowledge across those adapters to enhance various downstream tasks. In the future, we will evaluate our approach using some general domain KGs based on some general domain tasks. interprets 11
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
has method 12 has direct device 13 has dose form 14 has subject relationship context 15 has pathological process 16 has interpretation 17 moved to 18
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
has intent 19 has temporal context 
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
We evaluate our MoP on six datasets over various downstream tasks, including four question answering (i.e., PubMedQA, Jin et al. 2019;BioAsq7b, Nentidis et al. 2019;BioAsq8b, Nentidis et al. 2020;MedQA, Jin et al. 2020), one document classification (HoC, Baker and Korhonen 2017), and one natural language inference (MedNLI, Romanov and Shivade 2018) datasets. While HoC is a multi-label classification, and MedQA is a multichoice prediction, the rest can be formulated as a binary/multiclass classification tasks.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
• HoC (Baker and Korhonen, 2017): The Hallmarks of Cancer corpus was extracted from 1852 PubMed publication abstracts by Baker and Korhonen (2017), and the class labels were manually annotated by experts according to the Hallmarks of Cancer taxonomy. The taxonomy consists of 37 classes in a hierarchy, but in this paper we only consider the ten top-level ones. We use the publicly available train/dev/test split created by (Gu et al., 2020) and report the average performance over five runs by the average micro F1 across the ten cancer hallmarks.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
• PubMedQA (Jin et al., 2019): This is a question answering dataset that contains a set of research questions, each with a reference text from a PubMed abstract as well as an annotated label of whether the text contains the answer to the research question (yes/maybe/no ). We use the original train/dev/test split with 450/50/500 questions, respectively. The reported performance are the average of ten runs under the accuracy metric.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
• BioASQ7b, BioASQ8b (Nentidis et al., 2019(Nentidis et al., , 2020: The both BioASQ datasets are yes/no question answering tasks annotated by biomedical experts. Each question is paired with a reference text containing multiple sentences from a PubMed abstract and a yes/no answer. We use the official train/dev/test splits, i.e. 670/75/140 and 729/152/152 for BioASQ7b and BioASQ8b respectively, and the reported performances are the average of ten runs under the accuracy metric.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
• MedNLI (Romanov and Shivade, 2018): MedNLI is a Natural Language Inference (NLI) collection of sentence pairs extracted from MIMIC-III, a large clinical database. The objective of the NLI task is to determine if a given hypothesis can be inferred from a given premise. This task is formulated as the document classification task over three labels: {entailment, contradiction,neutral}. We use the same train/dev/test split generated by Romanov and Shivade (2018), and report the average accuracy performance over three runs.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
• MedQA (Jin et al., 2020): MedQA is a publicly available large-scale multiple-choice question answering dataset extracted from the professional medical board exams. It covers three languages: English, simplified Chinese, and traditional Chinese, but in this paper we only adopt the English set, which is split by Jin et al. (2020). Following (Jin et al., 2020), we use the Elasticsearch system to retrieve the top 25 sentences to each question+choice pair as the context for each choice, and concatenate them to obtain the normalized log probability over the five choices. Since this dataset is very large, we only report the average accuracy performance under three runs for all the models.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
Our MoP approach first infuses the factual knowledge for all the partitioned sub-graphs using their respective adapters with newly initialized parameters, then these knowledge-encapsulated adapters are further fine-tuned alone with the underlying BERT model through mixture layers. In this paper, we explored three approaches for implementing the mixture layers, which are described as follows:
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
• Softmax. As the default mixture layers deployed in our MoP, AdapterFusion is a recent proposed model (Pfeiffer et al., 2020a) that learns to combine the information from a set of tasks adapters by a softmax attention layer. In particular, the outputs from different adapters at layer l are combined using a contextual mixture weight calculated by a softmax over these adapters:
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
(3) For brevity, we denote our MoP with the original AdapterFusion mixture layers as Softmax.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
• Gumbel. We also extend the AdapterFusion by replacing the softmax layer with the Gumbel-Softmax (Jang et al., 2017) layer for obtaining more discrete mixture weights:
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
where g 1 , · · · , g K are i.i.d samples drawn from Gumbel(0, 1) distribution, and τ is a hyper-parameter controlling the discreteness. For brevity, we denote our MoP with the Gumbel-Softmax AdapterFusion mixture layers as Gumbel.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
• MoE. Mix-of-Experts (MoE) is a type of general purpose neural network component for selecting a combination of the experts to process each input. In particular, we use the the sparsely-gated mixture-of-experts, introduced by Shazeer et al. (2017), for obtaining a top-K sparse mixture of these adapters. And the mixture weights are calculated by:  where H(Φ l,G k ) is a function for transferring hidden variables into scalars with tunable Gaussian noise, and TopK(·) is a function for keeping only the top K values. We denote our MoP with this mixture approach as MoE. Table 6 shows the performance comparison of the three mixture approaches on the BioASQ7b and PubMedQA tasks. Note that Gumbel and MoE have additional hyper-parameters, i.e. τ and K, for controlling the discreteness and topK respectively. From Table 6, we can see that the original Adapter-Fusion with the softmax layer outperforms other two mixture approaches on all the hyper-parameter choices. This result justifies the choice of the mixture layers in our MoP model.
Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT
To further validate that the performance improvements of the evaluated BERTs using our MoP are gained due to the infused knowledge from the subgraph adapters, rather than the newly added more parameters of adapters, we split the partitioned sub-graphs into two groups according to their test performance ranking, and use our MoP to fine-tune the adapters of each group. Table 7 and 8 show the performances of all the adapters and our MoP combining the grouped adapters over train/dev/test sets of the BioASQ7b and PubMedQA datasets respectively. As we can see from the two tables, our MoP fine-tuned under the group of higher performance adapters can consistently obtain better performance than the group of the lower performance adapters. Note that we have shown that in Figure 2 
LM-CORE: Language Models with Contextually Relevant External Knowledge
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 ) store significant world knowledge in their parameters (Petroni et al., 2019).
LM-CORE: Language Models with Contextually Relevant External Knowledge
PLMs are typically fed a large amount of unstructured text which leads to the linguistic nuiances 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 . 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.").
LM-CORE: Language Models with Contextually Relevant External Knowledge
Large scale structured knowledge bases (KBs) such as YAGO (Suchanek et al., 2007) and Wikidata (Vrandečić 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 Zhang et al., 2019;Sun et al., 2020;Liu et al., 2020;Wang et al., 2021a,b Figure 1: Language Model Pre-Training with Contextually Relevant External Knowledge: 1 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). 2 An NER tagger is then applied to the masked input sequence (x) to identify named entities (underlined in black). 3 For the identified entities, the Knowledge Retrieval module fetches the set T x of all the triples from the Knowledge Base and then 4 scores all the retrieved triples using input-triple and input-relation similarity (details in Section 3.2). 5 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.
LM-CORE: Language Models with Contextually Relevant External Knowledge
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.
LM-CORE: Language Models with Contextually Relevant External 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).
LM-CORE: Language Models with Contextually Relevant External Knowledge
Task setting and Overview: Consider a language model L (such as BERT and RoBERTa) and a knowledge base K = {t hrt =< h, r, t > |h, t ∈ E; r ∈ R}. Here, we consider the knowledge base K as a set of triples such that each triple t hrt represents the relationship r between entities h and t. E is the set of all the entities, and 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 T x ∈ K such that the triples in 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 T x and is trained to make predictions using this additional knowledge.
LM-CORE: Language Models with Contextually Relevant External 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.
LM-CORE: Language Models with Contextually Relevant External Knowledge
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.
LM-CORE: Language Models with Contextually Relevant External Knowledge
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.
LM-CORE: Language Models with Contextually Relevant External Knowledge
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 Buf-fet 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.
LM-CORE: Language Models with Contextually Relevant External Knowledge
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.
LM-CORE: Language Models with Contextually Relevant External Knowledge
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 T x in Figure 1).
LM-CORE: Language Models with Contextually Relevant External Knowledge
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.
LM-CORE: Language Models with Contextually Relevant External Knowledge
Here, Emb(·) 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 (War-ren_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.
LM-CORE: Language Models with Contextually Relevant External Knowledge
where 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.
LM-CORE: Language Models with Contextually Relevant External Knowledge
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.
LM-CORE: Language Models with Contextually Relevant External 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.
LM-CORE: Language Models with Contextually Relevant External Knowledge
tokens in x and ind m is the index of the m th masked token.
LM-CORE: Language Models with Contextually Relevant External Knowledge
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.
LM-CORE: Language Models with Contextually Relevant External Knowledge
Data Sources and Pre-processing: We create our pre-training corpus using the December 20, 2018 snapshot of English Wikipedia that contains about 5.5M documents. Processing all the articles following the masking strategy described in Section 3.1 resulted in a total of ∼46.3M sentences with valid masks, from which we randomly sample sentences to create input examples.
LM-CORE: Language Models with Contextually Relevant External Knowledge
In order to illustrate the general nature of LM-CORE, we used two different PLMs as our Table 1: Mean precision at one (P@1) of various models on LAMA probe. We group all the models based on the base language model used (BERT or RoBERTa). For LM-CORE, (·, ·) 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  (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).
LM-CORE: Language Models with Contextually Relevant External Knowledge
We now present an analysis of how much, and if, having access to external knowledge can help PLMs in knowledge-intensive tasks. A popular way of assessing a model's ability to perform at such tasks is by using benchmark knowledge probes. We use the LAMA probe (Petroni et al., 2019) (Wang et al., 2021b) and CoLAKE (Sun et al., 2020) (based on RoBERTa) as the representative knowledge enhanced language models. Both KEPLER and Co-LAKE have used Wikidata5M as the knowledge base. We used author provided code and checkpoints for obtaining the reported numbers. For LM-CORE, we use four variants with different knowledge base and language encoder combinations as described above.
LM-CORE: Language Models with Contextually Relevant External Knowledge
We observe that our approach of providing external knowledge to the PLMs leads to substantially improved performance over the base language models and their SoTA knowledge enhanced variants. LM-CORE(b,w) achieves P@1 of 42.83% compared to 25.44% for BERT-large. Likewise, LM-CORE(r,w) achieves a P@1 of 41.96% significantly outperforming RoBERTa-large (24.24%). We also report the numbers on the four different subsets of LAMA revealing interesting insights. For all the models considered, we note that the performance on T-REx subset is higher than the Google-RE subset. We attribute this to the nature of knowledge required for probes in the four subsets. Note especially the column for Date of Birth (DoB) in the Table. All the models, except for LM-CORE(b,y) and LM-CORE(r,y) perform extremely poorly. This is because the Wikidata5M KB does not have date entity type and hence, the poor performance of models using Wikidata. We also note 
LM-CORE: Language Models with Contextually Relevant External Knowledge
PLMs are often sensitive to the linguistic variations in the input and are overly reliant on the surface form of entity names for making its predictions Poerner et al. (2020). For example, BERT can predict that a person with an Italian-sounding name was born in Italy even if this is factually incorrect. In order to evaluate the sensitivity and robustness of different models, we report the P@1 numbers for the LAMA-UHN (UnHelpfulNames) probing benchmark (Table 2) -a much harder subset of LAMA where input probes with helpful entity names are removed and the PLM has little or no helpful contextual signals from other tokens in the probe. We observe that the LM-CORE variants significantly outperform the base language models and their knowledge enhanced variants. Further, note that while all the baseline models suffer a significant fall in performance (expected due to the hardness of LAMA-UHN), the drop in performance of LM-CORE variants is much less. This indicates that having access to relevant external knowledge helps reduce the dependence on linguistic signals and results in the robust outperformance of LM-CORE variants.
LM-CORE: Language Models with Contextually Relevant External Knowledge
We analyze the effect of varying the number of candidates (k) during retrieval in Figure 3. We discuss with respect to Google-RE and T-REx subsets as our factual knowledge triples are most relevant for answering queries in these subsets (in comparison to commonsense queries in ConceptNet). We plot the Precision@1 (P@1) against increasing k values from 1 to 10 for LM-CORE(b,y) and LM-CORE(r,w) variants. We do not observe any consistent optimal k value across variants and data subsets. To add, there is no significant difference between P@1 values as k varies from 4 to 10. Hence, in order to maximize our recall while keeping the computational expense in mind, we select k = 8 for our experiments.
LM-CORE: Language Models with Contextually Relevant External Knowledge
We now study the role LM-CORE pre-training plays in helping the model access and utilize the retrieved knowledge and ensure that the model does not just rely on the knowledge stored in its parameters. We also study the effect of augmenting the base LMs with knowledge retrieved by LM-CORE. In addition to providing an insight into the quality of the knowledge retrieved by LM-CORE, this will also help us better understand the ability of LM-CORE to utilize the retrieved knowledge. We consider the following four variants on the LAMA probe (Table 3)  We observe that LM-CORE(r,w)'s performance (41.69 P@1) significantly exceeds RoBERTabase's performance using the same triples in input (30.06 P@1) , demonstrating that our training procedure equips the model with the capability of identifying and using relevant external knowledge effectively. There is a large drop in performance (from 41.69 P@1 to 19.51 P@1) when LM-CORE(r,w) is provided with random triples in input. This shows that the model exclusively accesses external knowledge to answer queries correctly. While the performance drops, it is important to note that the P@1 is similar to RoBERTa-base (20.46 P@1), highlighting that our training procedure does not lead to catastrophic forgetting and the model is able to rely on the knowledge stored in its parameters when semantically relevant triples are not provided in the input. Finally, although RoBERTa-base when augmented with contextually relevant triples does not perform competitively with LM-CORE, it demonstrates considerable improvement over the base RoBERTa model. This shows that high-quality relevant external knowledge has the potential to improve factual prediction, further reinforcing our motivation to train models to efficiently retrieve and use this knowledge.
LM-CORE: Language Models with Contextually Relevant External Knowledge
We consider two downstream tasks to study the effectiveness of LM-CORE for different NLP applications. We take Zero-Shot Relation Extraction (ZSRE) (Levy et al., 2017) and open-domain question answering over Web Questions (WQ) (Berant et al., 2013) dataset as the representative knowledge-intensive tasks. Tables 4 and 5 report the performance of LM-CORE and various other baselines for the two tasks, respectively. We use the LM-CORE(b,w) variant for these experiments as most baselines use BERT as the LM and Wikipedia as the knowledge base. For the ZSRE task, we use the data splits and evaluation systems provided as part of the KILT benchmark (Petroni et al., 2021). We find that for the ZSRE task, LM-CORE achieves a significantly higher F-1 score (74.80) compared to the second-best RAG model (49.95). Also, note that the online evaluator for the task considers exact string match (including casing, punctuations, etc.) for computing accuracy numbers but not for computing other metrics. Hence, the reported accuracy number for LM-CORE represent a lower bound as we don't have access to the same pre-processing pipeline to process its output. For the WQ dataset, we find that LM-CORE outperforms BERT with BM25 and neural retrievers, and the DrQA system. We observe that LM-CORE is outperformed by ORQA, designed explicitly for this task, and RAG (a retrieval augmented generative model). However, do note that all the models except LM-CORE have access to much larger knowledge source (complete Wikipedia corpus. ≈ 2B words), whereas LM-CORE only has access to the KB triples (21M triples, ≈ 140M words). As we show in the following subsection, with access to additional external knowledge, the performance of LM-CORE can improve significantly.  randomly sample 100 instances from the LAMA probe where the model failed and manually analyze these instances to identify the cases where the corresponding fact was not present in the YAGO KB. There were a total of of 41 such instances and we manually added the correct facts needed to answer the corresponding questions in YAGO. We then presented the 41 inputs again to the model with the updated KB. This time, the model used this newly added knowledge and was able to correct its prediction without any re-training for 36 out of 41 cases (87.8%). As discussed in the following sub-section (4.7), a majority of errors made by LM-CORE are due to missing facts in the KB and we expect that most of such errors can be corrected by having access to a larger, more comprehensive Knowledge Base.
LM-CORE: Language Models with Contextually Relevant External Knowledge
We now present some representative examples to illustrate the successes and failures of LM-CORE. Consider a test probe from the Google-RE subset of LAMA -Phil Mogg is a member of . Here, the correct output token is UFO, the band and BERT model incorrectly predicts parliament as the output token. This highlights the sensitivity of PLMs on context; BERT's prediction seems to be derived from its memorization of the frequently encountered phrase member of parliament during pretraining. We argue that the contextual knowledge retrieved by LM-CORE which includes the relevant fact <Phil Mogg; member of; UFO (band)> has helped the model to produce the correct output. We present more such successful examples in the Appendix (Tables 10 and 11).
LM-CORE: Language Models with Contextually Relevant External Knowledge
Next, we analyzed the cases where the proposed framework produced incorrect output and observed three major reasons for erros -(i) the required knowledge was not present in the knowledge base; (ii) the required knowledge was not retrieved despite being present in the knowledge base; and (iii) the system made errors after retrieving the rele-vant knowledge. The first problem cause could be addressed by enhancing the knowledge base as shown in Section 4.6. The other two causes of failure highlight the scope of improvement in our retrieval module as well as pre-training module, where further training could help the model make better use of the retrieved knowledge. Some representative examples of these different cases are presented in the Appendix (Table 12). Finally, we noticed some errors that could be attributed to the characteristics of the LAMA probe. Specifically, there are input probes that refer to entities without providing any additional context for disambiguation. For example, the sentence "James Johnson was born in " has no clues to determine whether the prompt is referring to the basketball player, Virginia congressman, or the Governor of Georgia with this name. We also noticed certain probes where there are multiple correct completions and the benchmark considers only one of these as the correct answer. For example, "Michelangelo is a by profession" can be correctly completed by poet, painter or architect, but the evaluation considers only poet as the correct answer. We also noticed some input examples with highly unambiguous language. For example, "X died in ", can refer to either X's place of death or date of death but only the former is accepted as the correct answer. Lastly, there are cases where slight (and correct) variations of the expected answer are evaluated as incorrect by the probe. For example, for the prompt "Harashima is citizen." Japan is provided as the correct answer while the prediction made by LM-CORE (Japanese) is considered incorrect.
LM-CORE: Language Models with Contextually Relevant External Knowledge
We presented LM-CORE, a framework to train language models with contextually relevant external knowledge. We show that having access to external knowledge leads to significant and robust outperformance over base language models and their knowledge enhanced versions on knowledge probing and two downstream tasks. We also showed how LM-CORE can handle knowledge updates and presented a thorough error analysis that helped us identify possible directions of future work. 
LM-CORE: Language Models with Contextually Relevant External Knowledge
We use the English Wikipedia (December 20, 2018) snapshot 3 to create our pre-training corpus and WikiExtractor 4 to process the dumps. This Wikipedia version contains about 5.5M documents. We retain the hyperlinks while extracting Wikipedia articles as we use them for creating entity masks (Section 3.1). Following the entity masking strategy described in Section 3.1, we obtain our pre-training corpus which contains ∼46.3M sentences in total. During pre-training the base LMs, we sample sentences containing valid masks. The pre-training corpus is maintained consistent across all LM-CORE variants.
LM-CORE: Language Models with Contextually Relevant External Knowledge
We use two popular knowledge bases (KBs) in LM-CORE -YAGO 4 (Suchanek et al., 2007) and Wikidata5M (Wang et al., 2021b). The statistics of the KBs -number of facts, entities and relations can be found in Table 6. We describe the preprocessing steps followed to obtain the respective final retrieval corpora in the following subsections.
LM-CORE: Language Models with Contextually Relevant External Knowledge
YAGO 4 5 is in RDFS format. YAGO facts are derived from Wikidata, however, all the entities are arranged in a taxonomy mapped to schema.org.
LM-CORE: Language Models with Contextually Relevant External Knowledge
We pre-process YAGO to remove triples involving relationships such as image, logo and url that point to meta-data such as images and other files. We also filter out triples that point to RDF literals or Wikidata URLs.
LM-CORE: Language Models with Contextually Relevant External Knowledge
We use the Wikidata5M subset of Wikidata as made available by Wang et al. (2021b) 6 . This subset of Wikidata is aligned with Wikipedia such that each 3 https://archive.org/ download/enwiki-20181220/ enwiki-20181220-pages-articles.xml.bz2 4 https://github.com/attardi/ wikiextractor entity in Wikidata5M has a corresponding entry in Wikipedia. We used the raw graph as provided in the dataset, the statistics of which are reported in Table 6. 
LM-CORE: Language Models with Contextually Relevant External Knowledge
We provide a closer look into our pre-training approach by showing examples of masked input sentences and the retrieved triple candidates from the knowledge base (Table 7). We observe that the facts retrieved are highly relevant for predicting the masked entities in the input context.
LM-CORE: Language Models with Contextually Relevant External Knowledge
B.1 How LM-CORE Compares with Other Retrieval Paradigms Table 8 also reports results of REALM (Guu et al., 2020) -a retrieval-based language model that retrieves relevant documents from a text corpus during pre-training. We observe that LM-CORE outperforms REALM on the ConceptNet, DoB (Google-RE), and 1-1 (T-REx) subsets, while REALM outperforms the proposed solution in other subsets of the LAMA probe. We specifically highlight an absolute 15 points improvement on the date-of-birth relation despite REALM using explicit date masks while training whereas our training corpus only has entity masks. This indicates that our model can use the the contextual knowledge provided by the retriever module even though it is not explicitly shown such knowledge during training. Note that while REALM is similar to our proposed solution as far as the idea of retrieving relevant knowledge is concerned, the key difference in the two approaches lies in the source of knowledge being used. REALM relies on an unstructured text corpus (Wikipedia) as the source of knowledge and employs a computationally complex retrieve and read paradigm requiring additional training of the knowledge retriever model. Our proposed solution, on the other hand, uses structured knowl-edge which offers the advantage of being (almost) unambiguous and less resource-hungry compared to unstructured text. We present the resource requirements of our approach and REALM in Table  9. Note that the size of the external knowledge (in number of words) used by REALM is an order of magnitude greater, and requires three times the number of parameters compared to our model. Furthermore, REALM was trained for 200K steps with a batch size of 512 on an 80 TPU cluster, whereas our proposed solution is much more efficient being trained for 1K steps with a batch size of 512 on a machine with 8 Nvidia A100 GPUs. This computational efficiency of our proposed solution allows us to continue further work on improving our performance by enhancing the structured knowledge base and bridge the performance gap with more complex and computationally expensive models such as REALM.
LM-CORE: Language Models with Contextually Relevant External Knowledge
We use the official LAMA data code 7 for evaluating P@1 numbers in Table 1. All the BERT-based models are evaluated using this repository. The LAMA code provides functionality for evaluating RoBERTa models trained in the fairseq framework. Hence, we evaluate RoBERTa-base, RoBERTalarge and KEPLER (Wang et al., 2021b) using this code. The KEPLER repo also points to this code for evaluation. CoLAKE (Sun et al., 2020), has adapted the official code 8 to allow huggingface transformer checkpoints as input, and hence this code is used for CoLAKE, LM-CORE(r,y) and LM-CORE(r,w) evaluation. We ensure the model vocabularies and data is consistent across evaluation. We have used author provided/recommended code and publicly available checkpoints from the official code repositories for all baselines.    
LM-CORE: Language Models with Contextually Relevant External Knowledge
We present various failure cases for LM-CORE in Table 12. These are representative of the type of errors we encountered, however, we observed that majority of the errors resulted due to correct facts missing from the KB.
LM-CORE: Language Models with Contextually Relevant External Knowledge
The complete LAMA-UHN results over all subsets of Google-RE and T-REx can be found in Table 13.
LM-CORE: Language Models with Contextually Relevant External Knowledge
We discuss the experimental setup and hyperparameter settings for our downstream tasks.
LM-CORE: Language Models with Contextually Relevant External Knowledge
We consider the open domain version of Zero Short RE (Levy et al., 2017) from Petroni et al. (2021). The dataset is split into three disjoint sets -train (147,909 samples, 84 relations), dev (3,724 samples, 12 relations) and test (4,966 samples, 24 re-lations). The systems are evaluated on relations never seen during training.
LM-CORE: Language Models with Contextually Relevant External Knowledge
We fine tune our model for 2 epochs with a batch size of 96. We use the Adam (Kingma and Ba, 2015) optimizer and a learning rate of 3e-5. We performed multiple trials by tuning the number of epochs in {1, 2, 5}.
LM-CORE: Language Models with Contextually Relevant External Knowledge
Web Questions (Berant et al., 2013) was created using questions that were sampled from the Google Suggest API. We used the same splits as Lee et al. (2019b) with training, dev and test sets containing 3417, 361 and 2032 samples respectively.
LM-CORE: Language Models with Contextually Relevant External Knowledge
We fine tuned our model for 20 epochs -we experimented with number of epochs in {10, 20, 30, 50}. We use the Adam (Kingma and Ba, 2015) optimizer and a learning rate of 3e-5.
LM-CORE: Language Models with Contextually Relevant External Knowledge
This work considers training of large language models using large textual corpora as well as structured knowledge bases. The model learns the nuances of the language and correlations between different real-world entities based on the data that is being used for training the model. Hence, there is a chance that the biases and noise in the training data will creep into the model parameters as well that can lead to a biased model behavior. We need to be careful in deploying the model and extrapolating the output of the model in applications such as search, conversational systems and recommendation systems where model's inherent biases can lead to catastrophic impacts on the user.
SKILL: Structured Knowledge Infusion for Large Language Models
Large pre-trained language models, such as BERT , GPT-3 (Brown et al., 2020), T5 (Raffel et al., 2019), REALM (Guu et al., 2020) and ERNIE  have become the state-of-the-art technology for many tasks. They are commonly pre-trained using unstructured text corpora, on tasks such as next word prediction, next sentence prediction (NSP) or masked language modelling (MLM). Especially for T5, selfsupervised learning on unlabelled text corpus with MLM has been a common pre-training recipe (Roberts et al., 2020). This is normally followed by a fine-tuning step on the task of interest (Ruder et al., 2019), although large language models have also proved useful without this task-specific finetuning (Brown et al., 2020).
SKILL: Structured Knowledge Infusion for Large Language Models
Beyond the capacity of contextual understanding, human-level language understanding pivots on the knowledge about the world. The world knowledge is often expressed as factual triples (c.f. Ji et al., 2020), in the form of (subject entity, relation, object entity). A knowledge graph (KG) defined by a set of factual triples consists of the subjects and objects as vertices/nodes, and the relations forming the edges connecting them. Most of the large scale KGs (e.g. Wikidata, Vrandečić and Krötzsch, 2014) are stored in triple format.
SKILL: Structured Knowledge Infusion for Large Language Models
LLMs demonstrate some capacity of learning world knowledge from the natural text corpus (Roberts et al., 2020), but it is unclear to what degree they are also able to learn and memorize new knowledge directly from structured KG triples, or from text describing them explicitly.
SKILL: Structured Knowledge Infusion for Large Language Models
In order to infuse knowledge into a LLM, one option is to generate a textual version of the knowledge base, and apply the standard training objectives, e.g. MLM. This is unfortunately highly nontrivial. One can either align sentences with KG triples, as done in ERNIE , or generate sentences from triples, as done in KELM (Agarwal et al., 2021). These approaches are unfortunately hard to port to knowledge graphs with different schemas. These processes are also lossy in that not every triple can be aligned or produce a valid sentence, and there is not a good understanding whether this can introduce unnecessary selection biases on top of biases existing in the original KG.
SKILL: Structured Knowledge Infusion for Large Language Models
In this work, we propose a method of Knowledge Infusion for Large Language Models (SKILL), where LLMs directly learns from knowledge triples. Experiment results shows the checkpoints trained with proposed method on Wikidata KG outperform the T5 baselines on four standard closed-book question-answering (QA) tasks. With a smaller KG, WikiMovies, the proposed method gain 3× exact match score performance improvement on MetaQA task. The models learning directly from knowledge triples performs competitively with the ones with the aligned natural sentences that contain the same amount of knowledge. Being able to learn directly from knowledge triples enables easy addition of structured knowledge into language modeling pre-training.
SKILL: Structured Knowledge Infusion for Large Language Models
Previous works that use knowledge graphs to enhance the quality of knowledge-intensive downstream tasks can be divided into two groups: using knowledge graphs at the inference time, and infusing knowledge into the model weights at the pre-training time. The proposed method falls in the latter group.
SKILL: Structured Knowledge Infusion for Large Language Models
Explicit usage of knowledge graphs. A retrieval-augmented model is commonly used, in order to retrieve and apply the knowledge from external memories or sources. FILM (Verga et al., 2021) and EaE (Févry et al., 2020) extend Transformer (Vaswani et al., 2017) models with external entity (both FILM and EaE) and fact (FILM) memories. REALM (Guu et al., 2020) is pre-trained to perform reasoning over a large textual knowledge corpus on-the-fly during inference. UniK-QA (Oguz et al., 2020) combines the structured and unstructured information to improve the open-domain QA tasks with a retriever-reader framework. The main difference between the proposed method, SKILL, and retrieval-augmented models is that SKILL doesn't introduce retrieval system or external memories to the model, but it directly embeds knowledge into the model parameters, which introduces no extra cost at inference time.
SKILL: Structured Knowledge Infusion for Large Language Models
Knowledge infusion. A common way of parameterized knowledge infusion is to map or convert structured knowledges into natural language text. ERNIE 3.0  trains a knowledgeenhanced model on a corpus combining triples and their aligned sentences, by randomly masking relation in a triple or words in a sentence. On the contrary, SKILL trains only on triples.
SKILL: Structured Knowledge Infusion for Large Language Models
KnowBert  incorporates knowledge from Wikipedia and WordNet (Miller, 1995) into a BERT model through entity embeddings with knowledge-attention and recontextualization mechanism. BERT-MK (He et al., 2020) is a BERT-based model that integrates graph contextual knowledge of a medical KG, which demonstrates the utility of graph-level knowledge. These approaches requires entity linking and sentences contextualizing the knowledge graph information.
SKILL: Structured Knowledge Infusion for Large Language Models
KG-FiD  extends the Fusion-in-Decoder model (Izacard and Grave, 2021) with a module that filters and re-ranks passages based on structural connections in knowledge graph between entities described in those passages. In contrast to the SKILL method that we propose, it requires the existence of natural text passages describing each knowledge graph entity, so Wikipedia corpus was used since it naturally provides articles that describe entities.
SKILL: Structured Knowledge Infusion for Large Language Models
Heinzerling and Inui (2021) explored the ability of language models to memorize and understand information from knowledge graphs, but used natural language representation of triples based on predefined templates instead of structured representation. Usage of predefined templates significantly limits scalability and therefore only relatively small knowledge graphs were used, such as Google-RE 1 .
SKILL: Structured Knowledge Infusion for Large Language Models
In contrast to the new method presented in this paper, all of these approaches require an explicit mapping between the knowledge graph entities or facts and corresponding natural language sentences, which can limit applications to industryscale knowledge graphs that don't have such a mapping.
SKILL: Structured Knowledge Infusion for Large Language Models
Different goals of using knowledge graphs. Besides that, some papers embed knowledge into model weights but pursue different goals rather than improving performance on downstream tasks. COMET (Bosselut et al., 2019) is most similar to our work and trains a commonsense-aware Transformer Language Model by learning to generate loosely structured commonsense descriptions in the natural language given the structured knowledge. Similar to us, it also uses KG triples in surface form as a source for training data, but in contrast to our research, the final goal of COMET is to generate new knowledge instead of utilizing existing ones. Another important difference is the scale: COMET uses Atomic  and Con-ceptNet (Speer et al., 2017) Knowledge Graphs that are much smaller than Wikidata (Vrandečić and Krötzsch, 2014).
SKILL: Structured Knowledge Infusion for Large Language Models
KELM (Agarwal et al., 2021) fine-tunes a T5 model to convert KGs to synthetic natural language sentences to augment existing pre-training corpora. We build our research on top of it and use the KELM dataset to compare structured and natural language representations of knowledge.
SKILL: Structured Knowledge Infusion for Large Language Models
There are two components of knowledge infusion for LLMs (SKILL): the corpus and the training method. We introduce the method based on Wikidata KG, but it can be applied to any other KGs.
SKILL: Structured Knowledge Infusion for Large Language Models
Training corpus. We use two corpora with different knowledge representations: Wikidata KG (Vrandečić and Krötzsch, 2014) in triple format, and KELM corpus 2 (Agarwal et al., 2021) as synthetic natural language sentences converted from Wikidata KG. The KELM corpus contains 15, 628, 486 synthetic sentences. To ensure two corpora share the same knowledge, we take the snapshot of the Wikidata KG used to created the KELM corpus, which contains 35, 697, 715 triples.
SKILL: Structured Knowledge Infusion for Large Language Models
To prevent the degradation of model performance on natural language understanding, we mix the Wikidata corpus or KELM corpus with natural text from C4 (Raffel et al., 2019), 50 : 50, for the knowledge infusion training data.
SKILL: Structured Knowledge Infusion for Large Language Models
Training method. T5 (Raffel et al., 2019) was trained through masked-language modelling with random span corruption on the C4 corpus. Roberts et al. (2020) found that masking salient terms (Guu et al., 2020) in pre-training T5 models, instead of masking random token spans, could significantly improve the performance on downstream tasks, e.g. closed-book QA.
SKILL: Structured Knowledge Infusion for Large Language Models
We apply salient span masking for unsupervised learning in our knowledge-infusing training. To mask the same amount of information is for both corpora, the following method is applied. For a knowledge triple, we mask either the subject or object entity. For a KELM sentence, we identify the aligned triple, with details in Appendix A, and mask the full spans corresponding to the subject or object in the triple. The relation tokens are never masked, as there is no robust way to map the abstract relation in knowledge triples to natural language tokens in KELM sentences. Examples of the inputs for both corpora are in Table 1.
SKILL: Structured Knowledge Infusion for Large Language Models
We assess SKILL by training and evaluating the knowledge infused models on closed-book QA tasks, where questions are provided without supporting context and external knowledge.
SKILL: Structured Knowledge Infusion for Large Language Models
SKILL pre-training. We apply SKILL on three T5.1.1 pre-trained checkpoints 3 , base, large, and XXL, with sizes of ∼ 250M, ∼ 800M and ∼ 11B parameters, respectively. For T5.1.1-base andlarge, SKILL training is performed for 500K steps with batch size 1024, which translates to ∼ 7.17 epochs on Wikidata KG and ∼ 16.38 epochs in KELM sentences. For T5.1.1-XXL, the model is trained for 100K steps to finish training in a feasible time.
SKILL: Structured Knowledge Infusion for Large Language Models
As baseline we use pre-trained T5 checkpoints of the same size. To make sure that improvements come from knowledge infusion instead of from longer C4 pre-training, we use a second baseline by further training the T5 checkpoints on C4 for half of the aforementioned steps, to match the amount of C4 pre-training used in SKILL.
SKILL: Structured Knowledge Infusion for Large Language Models
All the model variations are optimized by AdaFactor (Shazeer and Stern, 2018) with 10 −3 learning rate and 0.1 dropout rate, the same settings that were used for T5.
SKILL: Structured Knowledge Infusion for Large Language Models
Fine-tuning on closed-book QA tasks. We evaluate the checkpoints by fine-tuning on the following QA benchmarks: FreebaseQA (Jiang et al., 2019), WikiHop (Welbl et al., 2018), Triv-iaQA (Joshi et al., 2017) and NaturalQuestions (Kwiatkowski et al., 2019), with the aforementioned hyper-parameters for optimization and 128 batch size. For the benchmarks without a test split, we use the dev split for test, and the last 10% of train as dev split.
SKILL: Structured Knowledge Infusion for Large Language Models
The Exact Match (EM) scores on the test sets are calculated after being fine-tuned for 50K steps for T5.1.1-base and -large models, and 10K steps for -XXL models. All models converged with no noticeable over-fitting according to the EM scores on validation sets.
SKILL: Structured Knowledge Infusion for Large Language Models
Wikidata-answerable QA. We found that the majority of the questions in FreebaseQA and Wiki-   Hop can be answered directly from triples in Wikidata. This is because FreebaseQA was created by matching question-answer pairs with triples in Freebase (Bollacker et al., 2008), most of which was imported into Wikidata (Vrandečić and Krötzsch, 2014). For WikiHop, the questions were generated from Wikidata triples. However, TriviaQA and NaturalQuestions were created independently of Wikidata, and not every question can be answered using this knowledge base. We found frequent freshness issues, e.g. the golden answer for question "Who is the largest supermarket chain in the UK?" is "Aldi", while today it would be "Tesco". Some other questions can not be answered by WikiData, e.g. "Who, during a radio microphone test in 1984 said, 'I just signed legislation which outlaws Russia forever. The bombing begins in five minutes?'", with the golden answer "Ronald Reagan".
SKILL: Structured Knowledge Infusion for Large Language Models
To mitigate this, we created subsets of TriviaQA (TQA) and NaturalQuestions (NQ) that were somewhat more likely to have answers in Wikidata. We selected all the items for which there exist a triple in Wikidata that has the answer either as subject or object, and the other entity in the triple is mentioned in the question. We match the entities by entity name, case-insensitive. We name the Wikidataaligned version of TQA and NQ as TQA-matched and NQ-matched, respectively. The dataset sizes of all QA tasks are summarized in Appendix B.
SKILL: Structured Knowledge Infusion for Large Language Models
The results for closed-book QA tasks are summarized in Table 2. SKILL pre-trained models show improvements on FreebaseQA, WikiHop, and Wikidata-answerable versions of TriviaQA and Nat-uralQuestions, but no significant improvement on original TriviaQA and NaturalQuestions. As discussed in previous section, we believe this is due to the misalignment between the datasets and Wikidata.
SKILL: Structured Knowledge Infusion for Large Language Models
Models pre-trained on Wikidata KG gives competitive results with ones on KELM sentences. It shows that the triple representation is as good as natural language representation, while being much easier to scale up for larger KG.
SKILL: Structured Knowledge Infusion for Large Language Models
For T5.1.1-base and -large, additional pretraining on C4 boosts performance in comparison to the original baseline. For T5.1.1-XXL, this additional pre-training leads to a performance regress. In (Raffel et al., 2019), it is mentioned that training on C4 for multiple times may reduce the performance of a T5 model. Impact of model size. As shown in Figure 1, SKILL pre-training introduces bigger improvements when applied on larger models. With more than 35M triples in Wikidata KG, it is harder for    Table 3 demonstrate the effectiveness of SKILL pre-training, when it's possible to memorize the whole knowledge graph.
SKILL: Structured Knowledge Infusion for Large Language Models
As 1-hop questions are supported by single triples in the WikiMovies KG, a 3× improvement on EM score is observed for the sub-task. In order to answer 2/3-hop questions it is not enough to memorize the triples, the model needs to be able to reason with them. This requires a better understanding of the graph structure. Training with single triples may not be enough, and the observed improvement is notably smaller. The performance could be further improved by representing more explicitly the graph structure in the training data, which we leave for future work.
SKILL: Structured Knowledge Infusion for Large Language Models
We proposed a method to directly infuse knowledge from knowledge graphs into T5 models through pre-training. Empirical results show that T5 can learn directly from structured data and apply the learned knowledge to improve closed-book QA results. We also demonstrated that the models pre-trained on factual triples perform competitively with the ones on natural language sentences that contain the same knowledge. By enabling knowledge infusion directly from triples, this method can be very easily applied to industry-scale KGs.
SKILL: Structured Knowledge Infusion for Large Language Models
In this work, we are introducing a new method to pre-train a well known natural language understanding model, T5, on the full corpora of public knowledge graphs. To the best of our knowledge, the method will not introduce extra bias to either the model or the dataset beyond the one potentially inherited from Wikidata (Vrandečić and Krötzsch, 2014) and WikiMovies (Miller et al., 2016) knowledge graphs. On the other hand, through knowledge fusion pre-training introduced in this work, a language model will be able to learn factual information to improve the quality of parameterized knowledge embedded in the model, which is demonstrated by improvements on various closedbook question-answering tasks. The proposed method and recipe will provide positive impact to the natural language processing community and help to improve the factualness in pre-trained large language model checkpoints.
SKILL: Structured Knowledge Infusion for Large Language Models
A factual triple is the basic ingredient of a knowledge graph. However, as a semantic network, the graph structure of a knowledge graph describes how the factual triples are connected. This information is not easy to directly represent by random set of triples. We leave the exploration of how to infuse the information implied by the graph structure for future work. We expect that this will further improve the results, especially for multi-hop question-answering tasks.
SKILL: Structured Knowledge Infusion for Large Language Models
To find Wikidata KG entities in corresponding KELM sentences, we use Algorithm 1. Additional cycle on line 22 is needed because some entities have an information in brackets that should not be in a sentence, for example John Doe (born 1990). This algorithm matched at least one entity to 15, 383, 248 out of 15, 628, 486 KELM sentences. We don't try to match relation part of triples, because it could be represented in many different forms. For example, the triple (Pulp Fiction, cast member, John Travolta) could be represented as "John Travolta was an actor in Pulp Fiction", "John Travolta starred in Pulp Fiction", "John Travolta played Vincent Vega in Pulp Fiction", etc., and there is no way to robustly align a relation to all possible surface forms.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Recent years have witnessed the prominence of Pretrained Language Models (PLMs) (Peters et al., 2018;Radford et al., 2018;Devlin et al., 2019;Raffel et al., 2020;Xu et al., 2021) due to their superior performance on a wide range of languagerelated downstream tasks such as text classification (Kowsari et al., 2019), question answering (Rajpurkar et al., 2016), and machine reading comprehension (Nguyen et al., 2016). To fathom the prin-ciples of such effectiveness of PLMs, researchers have conducted extensive studies and suggested that PLMs have obtained rich knowledge during pre-training (Petroni et al., 2019;Davison et al., 2019). Hence, how to stimulate and exploit such knowledge is receiving increasing attention.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
One conventional approach to achieve that is fine-tuning (Devlin et al., 2019), where we add extra classifiers on the top of PLMs and further train the models under classification objectives. Fine-tuning has achieved satisfying results on supervised tasks. However, since the extra classifier requires adequate training instances to tune, it is still challenging to apply fine-tuning in few-shot learning (Brown et al., 2020) and zero-shot learning (Yin et al., 2019) scenarios. Originated from GPT-3 (Brown et al., 2020) and LAMA (Petroni et al., 2019(Petroni et al., , 2020, a series of studies using prompts (Schick and Schütze, 2021a; for model tuning bridge the gap between pre-training objective and down-stream tasks, and demonstrate that such discrete or continuous prompts induce better performances for PLMs on few-shot and zero-shot tasks.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
A typical way to use prompts is to wrap the input sentence into a natural language template and let the PLM conduct masked language modeling. For instance, to classify the topic of a sentence x: "What's the relation between speed and acceleration?" into the "SCIENCE" category, we wrap it into a template: "A [MASK] question: x". The prediction is made based on the probability that the word "science" is filled in the "[MASK]" token. The mapping from label words (e.g., "science" ) to the specific class (e.g., class SCIENCE) is called the verbalizer (Schick and Schütze, 2021a), which bridges a projection between the vocabulary and the label space and has a great influence on the performance of classification (Gao et al., 2021).
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Most existing works use manual verbalizers (Schick and Schütze, 2021a,b), in which the designers manually think up a single word to indicate each class. To ease the human effort of designing the class name, some works propose to learn the label words using discrete search (Schick et al., 2020) or gradient descent Hambardzumyan et al., 2021). However, the learnedfrom-scratch verbalizer, lack of human prior knowledge, is still considerably inferior to the manual verbalizers (see Appendix A for pilot experiments), especially in few-shot setting, and even not applicable in zero-shot setting, which leaves the manual verbalizer a decent choice in many cases.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
However, manual verbalizers usually determine the predictions based on limited information. For instance, in the above example, the mapping {science}→ SCIENCE means that only predicting the word "science" for the [MASK] token is regarded as correct during inference, regardless of the predictions on other relevant words such as "physics" and "maths", which are also informative. Such handcrafted one-one mapping limits the coverage of label words, thus lacking enough information for prediction and introducing bias into the verbalizer. Therefore, manual verbalizers are hard to be optimal in text classification, where the semantics of label words are crucial for predictions.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
The optimization-based expansion, though can be combined with manual verbalizers to yield better performance, only induces a few words or embeddings that are close to the class name in terms of word sense or embedding distance. Thus they are difficult to infer words across granularities (e.g. from "science" to "physics"). If we can expand the verbalizer of the above example into {science, physics} → SCIENCE, the probability of making correct predictions will be considerably enhanced. Therefore, to improve the coverage and reduce the bias of the manual verbalizer, we present to incorporate external knowledge into the verbalizers to facilitate prompt-tuning, namely, knowledgeable prompt-tuning (KPT). Since our expansion is not based on optimization, it will also be more favorable for zero-shot learning.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Specifically, KPT contains three steps: construction, refinement, and utilization. (1) Firstly, in the construction stage, we use external KBs to generate a set of label words for each label (in § 3.2). Note that the expanded label words are not simply synonyms of each other, but cover different granularities and perspectives, thus are more comprehensive and unbiased than the class name. (2) Secondly, to cope with the noise in the unsupervised expansion of label words, we propose four refinement methods, namely, frequency refinement, relevance refinement, contextualized calibration, and learnable refinement (in § 3.3), whose effectiveness is studied thoroughly in § 4. (3) Finally, we apply either a vanilla average loss function or a weighted average loss function for the utilization of expanded verbalizers, which map the scores on a set of label words to the scores of the labels.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
We conduct extensive experiments on zero-shot and few-shot text classification tasks. The empirical results show that KPT can reduce the error rate of classification by 16%, 18%, 10%, 7% on average in 0, 1, 5, 10 shot experiments, respectively, which shows the effectiveness of KPT. In addition to the performance boost, KPT also reduces the prediction variances consistently in few-shot experiments and yields more stable performances.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
In this section, we present our methods to incorporate external knowledge into a prompt verbalizer. We first introduce the overall paradigm of prompttuning and then elucidate how to construct, refine and utilize the knowledgeable prompt.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Let M be a language model pre-trained on large scale corpora. In text classification task, an input sequence x = (x 0 , x 1 , ..., x n ) is classified into a class label y ∈ Y. Prompt-tuning formalizes the classification task into a masked language modeling problem. Specifically, prompt-tuning wraps the input sequence with a template, which is a piece of natural language text. For example, assuming we need to classify the sentence x ="What's the relation between speed and acceleration?" into label SCIENCE (labeled as 1) or SPORTS (labeled as 2), we wrap it into
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Then M gives the probability of each word v in the vocabulary being filled in [MASK] token P M ([MASK] = v|x p ). To map the probabilities of words into the probabilities of labels, we define a verbalizer as a mapping f from a few words in the vocabulary, which form the label word set V, to the label space Y, i.e., f : V → Y. We use V y to denote the subset of V that is mapped into a specific label y, ∪ y∈Y V y = V. Then the probability of label y, i.e., P (y|x p ), is calculated as
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
where g is a function transforming the probability of label words into the probability of the label. In the above example, regular prompt-tuning may define V 1 = {"science"}, V 2 = {"sports"} and g as an identity function, then if the probability of "science" is larger than "sports", we classify the instance into SCIENCE. We propose KPT, which mainly focuses on using external knowledge to improve verbalizers in prompt-tuning. In KPT , we use KBs to generate multiple label words related to each class y, e.g., V 1 = {"science","physics", ...}. And we propose four refinement methods to eliminate the noise in the expanded V. Finally, we explore the vanilla average and weighted average approaches for the utilization of the expanded V. The details are in the following sections.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
The process of predicting masked words based on the context is not a single-choice procedure, that is, there is no standard correct answer, but abundant words may fit this context. Therefore, the label words mapped by a verbalizer should be equipped by two attributes: wide coverage and little subjective bias. Such a comprehensive projection is crucial to the imitation of pre-training, which is the essence of prompt-tuning. Fortunately, external structured knowledge could simultaneously meet both requirements. In this section, we introduce how we use external knowledge for two text classification tasks: topic classification and sentiment classification.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
For topic classification, the core issue is to extract label words related to the topic from all aspects and granularities. From this perspective, we choose Related Words 1 , a knowledge graph G aggregated from multiple resources, including word embeddings, ConceptNet (Speer et al., 2017), WordNet (Pedersen et al., 2004), etc., as our external KB. The edges denote "relevance" relations and are annotated with relevance scores. We presume the the name of each class v 0 is correct and use them as the anchor node to get the neighborhood nodes N G (v 0 ) whose scores are larger than a threshold η as the related words 2 . Thus, each class is mapped into a set of label words V y = N G (v 0 ) ∪ {v 0 }. For binary sentiment classification, the primary goal is to extend the binary sentiment to sentiment of more granualities and aspects. We use the sentiment dictionary summarized by previous researchers 3,4 . Several examples of the label words in the KPT are in Table 1.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Although we have constructed a knowledgeable verbalizer that contains comprehensive label words, the collected label words can be very noisy since the vocabulary of the KB is not tailored for the PLM. Thus it is necessary to refine such verbalizer by retaining high-quality words. In this section, we propose four refinement methods addressing different problems of the noisy label words.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Frequency Refinement. The first problem is to handle the rare words. We assume that several words in the KB are rare to the PLM, thus the prediction probabilities on these words tend to be inaccurate. Instead of using a word-frequency dictionary, we propose to use contextualized prior of the label words to remove these words. Specifically, given a text classification task, we denote the distribution of the sentences x in the corpus as D.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
For each sentence in the distribution, we wrap it into the template and calculate the predicted probability for each label word v in the masked position P M ([MASK]= v|x p ). By taking the expectation of the probability over the entire distribution of sentences, we can get the prior distribution of the label words in the masked position. We formalize it as
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Empirically, we found that using a small-size unlabeled support setC sampled from the training set and with labels removed, will yield a satisfying estimate of the above expectation. Thus, assuming that the input samples {x ∈C} have a uniform prior distribution, the contextualized prior is approximated by
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Then we remove the label words whose prior probabilities are less than a threshold. Details can be found in Appendix C. Relevance Refinement. As our construction of knowledgeable label words is fully unsupervised, some label words may be more relevant to their belonging class than the others. To measure the relevance of a label word to each class, we obtain the prediction probability of the label word on the support setC as the vector representation q v of the label words, i.e., q v 's i-th element is
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
where x ip represents the sentence x i combined with the template p.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
To estimate the class's representation, we presume that the name of each class v 0 , such as "science" for SCIENCE, though lack of coverage, is very relevant to the class. Then we use the vector representation q v 0 of the these names as the class's representation q y . Therefore the relevance score between a label word v and a class y is calculated as the cosine similarity between the two representation:
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Moreover, some label words may contribute positively to multiple classes, resulting in confusion between classes. For example, the potential label word "physiology" of class SCIENCE may also be assigned with a high probability in a sentence of class SPORTS. To mitigate such confusion and filter the less relevant label words, we design a metric that favors the label word with high relevance merely to its belonging class and low relevance to other classes:
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
where f (v) is the corresponding class of v. Ideally, a good label word should at least has a higher relevance score for its belonging class than the average relevance score for the other classes. Therefore, we remove the label words with R(v) < 1. In practice, we have a slight modification to Equation (6), please refer to appendix C for details.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Essentially, this Relevance Refinement adopts the idea of the classical TF-IDF (Jones, 1972) algorithm which estimates the relevance of a word to a document. It prefers to use a word that is relevant to a specific document while irrelevant to other documents as the keyword of the document. In KPT, a class is analogous to a document, while a label word is comparable to the word in the document. From this perspective, equation (6) is a variant of TF-IDF metric.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Contextualized Calibration. The third problem is the drastic difference in the prior probabilities of label words. As previous works Holtzman et al., 2021) have shown, some label words are less likely to be predicted than the others, regardless of the label of input sentences, resulting in a biased prediction. In our setting, the label words in the KB tend to have more diverse prior probabilities, resulting in a severer problem (see Table 2). Therefore, we use the contextualized prior of label words to calibrate the predicted distribution, namely, contextualized calibration (CC):
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
where P D (v) is the prior probability of the label word. The final probability is normalized to 1. Learnable Refinement. In few-shot learning, the refinement can be strengthen by a learning process. Specifically we assign a learnable weight w v to each label word v (may be already refined by the previous methods). The weights form a vector w ∈ R |V| , which is initialized to be a zero vector. The weights are normalized within each V y :
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Intuitively, in the training process, a small weight is expected to be learned for a noisy label word to minimize its influence on the prediction. Note that in few-shot setting, calibration may not be necessary because the probability of a label word can be trained to the desired magnitude, i.e.,
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
In addition to these refinement methods, since many label words are out-of-vocabulary for the PLM and are split into multiple tokens by the tokenizer. For these words, we simply use the average prediction score of each token as the prediction score for the word. The influence of this simple approach is studied in Appendix D.3.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
The final problem is how to map the predicted probability on each refined label word to the decision of the class label y.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Average. After refinement, we can assume that each label word of a class contributes equally to predicting the label. Therefore, we use the average of the predicted scores on V y as the predicted score for label y. The predicted labelŷ iŝ
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
We use this method in zero-shot learning since there is no parameter to be trained.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Weighted Average. In few-shot setting, supported by the Learnable Refinement, we adopt a weighted average of label words' scores as the prediction score. The refinement weights α v are used as the weights for averaging. Thus, the predictedŷ isŷ
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
where s(y|x p ) is
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
This objective function is suitable for continuous optimization by applying a cross-entropy loss on the predicted probability.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
We provide a theoretical illustration of the KPT framwork in Appendix B.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
We evaluate KPT on five text classification datasets to demonstrate the effectiveness of incorporating external knowledge into prompt-tuning.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
We carry out experiments on three topic classification datasets: AG's News (Zhang et al., 2015), DB-Pedia (Lehmann et al., 2015), and Yahoo (Zhang et al., 2015), and two sentiment classification datasets: IMDB (Maas et al., 2011) and Amazon (McAuley and Leskovec, 2013). The statistics of the datasets are shown in Table 7. The detailed information and the statistics of each dataset is in Appendix E. We test all prompt-based methods using four manual templates and report both the average results (with standard error) of the four templates and the results of the best template (shown in (brackets) ). The reasons for using manual templates and the specific templates for each dataset are in Appendix E.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Our experiments are based on OpenPrompt , which is an open-source toolkit to conduct prompt learning. For the PLM, we use RoBERTa large (Liu et al., 2019) for all experiments. For test metrics, we use Micro-F1 in all experiments. For all zero-shot experiments, we repeat the experiments 3 times using different random seeds if randomness is introduced in the experiments, and for all few-shot experiments, we repeat 5 times. Note that considering the four templates and five/three random seeds, each reported score of prompt-based methods is the average of 20/12 experiments, which greatly reduces the randomness of the evaluation results. For the refinement based on the support setC, the size of the unlabeled support set |C| is 200. For few-shot learning, we conduct 1, 5, 10, and 20-shot experiments. For a k-shot experiment, we sample k instances of each class from the original training set to form the fewshot training set and sample another k instances per class to form the validation set. We tune the entire model for 5 epochs and choose the checkpoint with the best validation performance to test. Other hyper-parameters can be found in Appendix F.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
In this subsection, we introduce the baselines we compare with. To better understand our proposed methods, we also compare within the performance of KPT using different configuration.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Fine-tuning (FT). Traditional fine-tuning method inputs the hidden embedding of [CLS] token of the PLM into the classification layer to make predictions. Note that fine-tuning can not be applied to the zero-shot setting, since the classification layer is randomly initialized.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Prompt-tuning (PT). The regular prompttuning method uses the class name as the only label word for each class, which is used in PET (Schick and Schütze, 2021a) and most existing works. For a fair comparison, we do not use the tricks in PET, such as self-training and prompt ensemble, which are orthogonal to our contributions.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Automatic Verbalizer (AUTO). The automatic verbalizer is proposed by PETAL (Schick et al., 2020), which uses labeled data to select the most informative label words inside a PLM's vocabulary. It is targeted at the situation when no manually defined class names are available. It's not obvious how to combine it with the manually de-   fined class name to boost the performance, and how it can be applied in a zero-shot setting. Therefore we only compare it in the few-shot setting with no class name information given.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Soft Verbalizer (SOFT). The soft verbalizer is proposed by WARP (Hambardzumyan et al., 2021). They use a continuous vector for each class and use the dot product between the masked language model output and the class vector to produce the probability for each class. In our experiments, its class vectors are initialized with the class names' word embedding, since it is more effective with manual class names as the initial values (see Appendix A). As an optimization-based method, Soft Verbalizer is not applicable in the zero-shot setting.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
PT+CC. For zero-shot setting, we further introduce PT combined with our proposed contextualized calibration 5 as a baseline to see how much improvement is made by contextualized calibration instead of knowledgeable verbalizers.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
For KPT , we experiment with different variants to better understand the proposed methods such as refinement. -FR, -RR, -CC and -LR is the variant that does not conduct Frequency Refinement, Relevance Refinement, Contextualized Calibration, and Learnable Refinement, respectively. In few-shot experiments, we presume that the supervised training data can train the output probability of each label word to the desired magnitude, thus we don't use CC and FR in the KPT . This decision is justified in Appendix D.2.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
In this subsection, we introduce the specific results and provide possible insights of KPT .
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Zero-shot. From Table 2, we see that all the variants of KPT , except for KPT-CC, consistently outperforms PT and PT+CC baselines, which indicates the effectiveness of our methods. Comparison between PT and PT+CC proves that Contextualized Calibration is very effective in the zero-shot setting. The results of KPT-FR-RR-CC, which is the variant without any refinement, reveal the label noise is severe in the automatically constructed knowledgeable label words. The gap between KPT-FR-RR and KPT-FR-RR-CC is larger than the gap between PT+CC and PT, demonstrating the drastic difference in the prior probabilities of the knowledgeable label words as we hypothesized in § 3.3. Comparison between KPT, KPT-FR, KPT-FR-RR proves the effectiveness of the refinement methods.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
For the analysis regarding each type of classification task, we observe that the performance boost compared to the baselines in topic classification is higher than sentiment classification, which we conjecture that topic classification requires more external knowledge than sentiment classification. While CC offers huge improvement (on average +13%) over PT baseline, the incorporation of external knowledge further improves over PT+CC up to 11% on DBPedia, and 6% on AG's News and Yahoo. We also observe that the improvement brought 5 The same support sets are used as KPT . by the refinement methods is more noticeable for topic classification tasks. By looking at the fraction of label words maintained after the refinement process (See appendix D.4), we conjecture that the sentiment dictionary that we used in sentiment classification tasks contains little noise. Moreover, the improvement brought by the refinement process justifies the resilience of our methods to recover from noisy label words.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Few-shot. From Table 3, we first find out that prompt-based methods win over fine-tuning by a dramatic margin under nearly all situations. The gap enlarges as the shot becomes fewer. Comparing the baseline methods, the Soft Verbalizer (SOFT) generally wins over the Manual Verbalizer(PT) by a slight margin. However, automatic verbalizer (AUTO), although free of manual effort, lags behind the other verbalizers especially in a low-shot setting. The reason is obvious since the selection of label words among the vocabulary becomes inaccurate when labeled data is limited.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
When comparing KPT with the baseline methods, we find KPT or its variants consistently outperform all baseline methods. On average, 17.8% , 10.3%, and 7.4% error rate reduction from the best baseline methods are achieved on 1, 5, and 10 shot experiments, respectively. Comparing within the variants of KPT , we find that RR and LR are generally effective across shots on topic classification dataset, while in sentiment classification dataset, KPT works well without the refinements, which is consistent with our previous assumptions that the sentiment dictionary has little noise. Note that the KPT-RR variant does not utilize any unlabeled support setC since we do not conduct CC and FR by default in few-shot learning. This variant is still superior to the baseline methods in most cases. In terms of variance, we can see that KPT enjoys smaller variances than baseline methods in most cases, demonstrating that the better coverage of label words stabilizes the training.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
For 20-shot experiments, we can see that the gap between different methods narrows as the training data becomes sufficient. However, KPT and its variants still win by a consistent margin over the baseline methods. Surprisingly, with more training data, LR does not become more powerful as we may hypothesize. We conjecture that it is because all label words, even with some noise, can serve as training objectives for prompt tuning. This perspective is similar to Gao et al. (2021) that using "bad" as a label word for the class "positive" can still preform classification although the performance degrades.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Ablation studies about our refinement methods have been shown in the previous section. In this section and Appendix D, we conduct more in-depth analyses on the proposed methods.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
One advantage of KPT is that it can generate diverse label words across different granularities. To specifically quantify such diversity, we conduct a case study. For the correctly predicted sentences of a class y, we count the frequency of label words v ∈ V y appearing in the top-5 predictions for the [MASK] position. Then we report the top-15 frequent label words in Figure 2. Due to space limit, only the results of SPORTS and BUSINESS category of AG's News are shown. As shown in Figure 2, a diversity of label words, instead of mainly the original class names, are predicted. And the predicted label words cover various aspects of the corresponding topic. For example, for the topic SPORTS, the predicted "leagues", "football", and "coach" are related to it from different angles.  
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
In addition to the visualization, we study the influence of the support set's size on zero-shot text classification in Appendix D.1. Then we justify that few-shot learning via labeled data eases the need for calibration and frequency-based refinement in Appendix D.2. We also demonstrate that our approach to handling the out-of-vocabulary (OOV) words is reasonable in Appendix D.3. Moreover, we take a closer look at the refinement process by analyzing the fraction of label words maintained during refinement in Appendix D.4. Finally, we discuss the potential use of the proposed methods when knowledge bases resources are not readily available in Appendix D.5.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
In this paper, we propose KPT , which expands the verbalizer in prompt-tuning using the external KB. To better utilize the KB, we propose refinement methods for the knowledgeable verbalizer. The experiments show the potential of KPT in both zero-shot settings and few-shot settings. For future work, there are open questions related to our research for investigation: (1) Better approaches for combining KB and prompt-tuning in terms of template construction and verbalizer design.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
(2) Incorporating external knowledge into prompt-tuning for other tasks such as text generation.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Zhiyuan Liu, Huadong Wang and Shengding Hu proposed the idea and led the research. Shengding Hu designed the methods and conducted the experiments. Ning Ding and Shengding Hu wrote the abstract, introduction (Section 1) and method (Section 3) part of the paper. Shengding Hu finished the other parts. Ning Ding, Huadong Wang and Zhiyuan Liu thoroughly revised the whole paper. Jingang Wang, Juanzi Li, Wei Wu and Maosong Sun proofread the paper and provided valuable comments.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 12697-12706. PMLR.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
As pointed out by (Gao et al., 2021), manually defined verbalizer is competitive or even better than automatically searched/optimized verbalizers, which strengthens our motivation to improve over manual verbalizers by injecting more external human knowledge. To further illustrate the advantage of manual verbalizer, we conduct pilot experiments in soft verbalizer. Soft Verbalizer (Hambardzumyan et al., 2021) can be initialized with the predefined class name as the label words, which is adopted by us as a baseline in Table 3. It can also be randomly initialized without the manually defined class names. We test the performance of Soft Verbalizer with and without the manually defined class name in 5 and 10 shot experiments. From Table 4, we can see that the gaps between the variants are generally large. Therefore further improving the verbalizer with manually defined class name is a promising direction than the learned-from-scratch verbalizer without any human prior.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
In this section we provide a theoretical analysis of the whole framework used by KPT . In prompt tuning, given a text x, we wrap it into a template to form a wrapped sentence x p . We then predict the probability of the label word v using a PLM:
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
where [M] is short for [MASK], denoting the label word's prediction at the masked position of the template. Then, if multiple label words are used to contribute to a single label, the predicted probability of the label is defined by marginalizing the probability of predicting all the label words, i.e.,
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Since the prediction of Y is independent of x given v, we can write Equation (13) 
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Using the Bayes Theorem and assuming a balanced classification problem, Equation (14)    Now, the prediction probability of the label is composed of three parts.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
(1) The first part p(v|Y = y) is the probability of predicting the specific label word v given the class label y. Intuitively, if a label word is relevant to label y, this term will be assigned a high probability. In KPT , the Relevance Refinement estimate this probability using a quantized objective, i.e., if a relevance score exceeds the threshold 1, it will be maintained, otherwise, it will be filtered. On the other hand, Learnable Refinement estimates this probability using continuous weights.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
(2) The second part is p([M]= v) in the denominator. This term is actually the prior probability of label words v, which is estimated by our Contextualized Calibration. Previous works also try to approach this term using a context-free manner (Holtzman et al., 2021;.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
(3) The last term P M ([MASK]= v|x p ) is the probability of the label words v predicted by the PLM, which is the only component in most works such as Manual Verbalizers (Schick and Schütze, 2021a), yielding a sub-optimial solution compared to KPT .
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Verbalizers with multiple label words for a class label can all be formalized into this framework once they use Equation (13) as their backbone hypothesis. However, to the best of our knowledge, KPT is the first to combine all of the three components to form a powerful verbalizer.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
In this section, we detail the refinement process by making some practical modifications to the methods in § 3.3.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Frequency Refinement. For Frequency Refinement, since the absolute value distribution of the contextualized prior probability may be different for each task, determining a specific threshold of the contextualized prior probability may be tricky and elusive. We use a ranking-based threshold, i.e., we filter the label words that appear in the lower half of the contextualized prior probability.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Relevance Refinement. For Relevance Refinement, we observe that in the classification task with only a few classes, it's better to provide a stricter criterion to ensure that the relevance scores of a label word to any other class is lower than the score to the belonging class, i.e., maximum in the term of IDF-score is preferred. To keep a unified criterion, we use a norm-based IDF-score.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
This criterion will approximate the maximum value in {r(v, y)|y ∈ |Y|, y ̸ = f (v)} in classification with only a few labels, and revert to the mean score in Equation (6) when conducting classification with many labels. We take C = 10 (without trial and error) in the experiments. And 0 < ϵ ≪ 1 is a small number to prevent numerical error.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Calibration.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
In fine-tuning, calibration has been studied under the topic of prediction confidence and out-ofdistribution detection (Kong et al., 2020). Recently, it got renascent attention in prompt learning Holtzman et al., 2021). In prompt learning, the PLM has a natural tendency to predict one word over another word regardless of the real sentence input. For example, GPT-3 prefers to predict "positive" over "negative" given "N/A" as the input sentence . Therefore the calibration is crucial (see Table 2) when no posterior optimization is conducted, i.e., in zeroshot learning. Existing methods such as PMI DC propose only using the empty template without filling the template with the instances in the corpus, for example, "A [Mask] question :", to produce the calibration logits. Our proposed Contextualized Calibration utilizes a limited amount of unlabeled support data to yield significantly better results. However, since we target the data-scarce scenario, we study in detail the amount of unlabeled data necessary to produce a satisfying calibration result.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
In Figure 3, we draw the performance of KPT -RR with different support set sizes |C|. We also draw the performance of PMI DC on the |C| = 0 for comparison. From Figure 3, we find that |C| ∼ 50 is enough to yield a satisfying calibration. Contextualized calibrate is more effective in classification with many classes, while calibrate without the context is effective in classification with few classes.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
In addition, we must point out that if we have a set of sentences to classify in real-world scenarios, we can use these sentences themselves as the support set to conduct more accurate Contextualized Calibration.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Calibration.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Although calibration is crucial for the zero-shot setting, we do not perform calibration for the fewshot setting because we assume that the posterior probability of the label words can be trained to the desired magnitude with only a few training instances. We also do not perform Frequency Refinement for few-shot learning due to the same assumption. To verify the assumption empirically, we add both Contextualized Calibration and Frequency Refinement to KPT and test the performance under different settings. The results are shown in Table 5.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
The performance comparison to KPT without CC and FR in Table 3 are shown using up arrows and down arrows. We can see that except in Yahoo, the improvement is not consistent for even negative, which supports our assumption that the need for calibration is greatly eased with the supervised input data.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
Since the knowledgeable verbalizer is expanded using external resources which may not be tailored for the vocabulary of PLM. Thus, many label words are out-of-vocabulary (OOV) and are split into multiple tokens by the tokenizer. For these words, as mentioned in § 3.3, we average the prediction probability of each token in the single [MASK] position, which may not be very reasonable at the first glance. Therefore, we conduct an ablation study that whether forcing the label words to be a single token in the vocabulary of the PLM leads to better performance. The results under different shots are shown in Table 6. Surprisingly, making the singletoken restriction does not yield stable improvement, instead, in many cases, the performance degrades by minor margins. Therefore we conclude that our method to handle OOV label words that are split by the tokenizer into multiple tokens is simple yet reasonable. More importantly, the label words expanded by the knowledge bases but not in the PLM's vocabulary can serves as good label words in prompt tuning as well.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
In this section, we report the number of label words that remained after Frequency Refinement and Relevance Refinement process. As we can see, these refinement methods remove a large fraction of label words while retaining the ones that are most informative. However, even the fewest number of remaining label words exceeds 100, which is far more than the number of label words in the previous works (Schick et al., 2020). The broad coverage of label words contributes to the success of KPT .   Table 3, and the red down arrow  means the results is lower than KPT in Table 3 Table 6: Results of restricting the expanded label word to be a single token in the PLM's vocabulary, where ST denotes "single token". The green up arrow  means the results is higher than KPT in Table 3, and the red down arrow  means the results is lower than KPT in Table 3. Although KBs are ubiquitous in natural language processing, there are cases that no readily available KBs can be found for specific tasks. For these tasks, if we have enough unlabeled corpus, we can use the methods proposed by LOTClass (Meng et al., 2020) to mine potential label words from the corpus. More specifically, LOTClass (Meng et al., 2020) uses a self-supervised objective to train the PLM to extract the topic-related words from the whole unlabeled training corpus. Experiments that combine KPT with LOTClass are beyond the scope of our work, but we believe the combination of the two can be very effective.
Knowledgeable Prompt-tuning: Incorporating Knowledge into Prompt Verbalizer for Text Classification
We carry out experiments on three topic classification datasets: AG's News (Zhang et al., 2015), DB-Pedia (Lehmann et al., 2015), and Yahoo (Zhang et al., 2015), and two sentiment classification datasets: IMDB (Maas et al., 2011) and Amazon (McAuley and Leskovec, 2013). The statistics of the datasets are shown in Table 7.  Due to the rich expert knowledge contained, the manual templates are proven to be competitive with or better than auto-generated templates (Gao et al., 2021) even though they are simpler to be constructed. Therefore we use manual templates in our experiments. Manual templates are also more applicable than auto-generated templates in the zero-shot setting. To mitigate the influence of different templates, we test KPT under four manual templates that are either introduced by (Schick and Schütze, 2021a) or tailored to fit the dataset for each experimental configuration. The templates for each dataset is listed below.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Neural models have greatly advanced the state of the art in natural language processing and generation tasks. Accordingly, there has been increasing interest in applying neural language models to tasks which require extensive world knowledge to solve (Petroni et al., 2021). Much of this world knowledge can be found distributed over text corpora, which raises the question whether language models pre-trained on text corpora capture this information. Recent work suggests that while language models may successfully predict facts about the world (Petroni et al., 2019) such knowledge is superficial and unreliable . Our goal is to reliably incorporate information from across a text corpus into a language model. Recent work has represented the information present in a text corpus explicitly by constructing a virtual knowledge base (VKB) (Dhingra et al., 2020;. A VKB consists of dense representations of entity mentions in the text, designed to reflect the property or relation expressed by the entity mention. We propose to incorporate a VKB into a language model by using it as an external memory, performing attention over the entire VKB within a Transformer model. In this way the model can synthesise and reason over many disparate sources of information from the text corpus. We refer to the VKB used in such a way as Mention Memory, and the model as TOME (Transformer Over Mention Encodings). We first pre-train a mention encoder to specifically encourage mention representations that are useful for a Transformer model, and construct a Mention Memory from 150 million entity mentions in English Wikipedia. Then we train a TOME model with attention layers over the Mention Memory, which is kept frozen (see Figure 1).
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
We argue that the Mention Memory approach has several appealing properties. First, TOME retrieves entity mention representations corresponding to specific entity attributes or relations described in the Figure 1: Overview of Mention Memory. A pre-trained mention encoder is used to generate dense representations for each entity mention in Wikipedia (approximately 150 million total) which are stored in a table. The TOME model takes a passage annotated with entity mention boundaries as input, and applies a Transformer block. Next, the TOME model applies one or more TOMEBlocks. Each TOMEBlock contains a memory attention layer and a Transformer block.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
corpus. This retrieval is much more fine-grained than aggregate entity retrieval methods such as Entities as Experts (EaE) (Févry et al., 2020), and we show large improvements in accuracy over EaE on tasks that require detailed entity information, such as claim verification and entity-based question answering. The fine-grained retrieval also allows potential users to see more precisely what knowledge the model's predictions is based on (see Table 4). Second, TOME retrieves dense representations, which are easy to incorporate into a Transformer model without reprocessing the input, unlike raw text. Therefore, TOME is able to retrieve, assimilate and reason over information from many different sources within a single Transformer model, allowing for multi-source and multi-hop reasoning without the beam search machinery that is required for multi-hop retrieve-and-read (Zhao et al., 2021). This also makes TOME much more scalable: retrieve-and-read approaches have to read many retrieved passages which becomes expensive with larger reader models, while the cost of memory layers does not scale with reader size and is negligible for larger readers. Third, the retrieval is latent, without direct or distant supervision on the retrieved results. We show that, even without supervision, the model learns to retrieve highly specific and informative entity attributes and perform multiple reasoning steps. Finally, the memory table is semi-parametric, so knowledge can be added or updated by applying the mention encoder to new text without retraining.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
In order to verify the model's capacity to capture accurate factual information in the corpus, we start by evaluating TOME on the HoVer (Jiang et al., 2020), FEVER (Thorne et al., 2018) andFM2 (Eisenschlos et al., 2021) claim verification datasets, on which it strongly improves performance over entity aggregate and comparable retrieve-and-read baselines. We demonstrate that the model learns to attend to informative mentions for verifying claims using only the verification accuracy as a signal. Ablations show the memory is crucial for performance, and that the model can effectively use larger memory than it was pre-trained on. In a second set of experiments we evaluate TOME on question-answering benchmarks TriviaQA (Joshi et al., 2017), ComplexWebQuestions (Talmor & Berant, 2018) and EntityQuestions (Sciavolino et al., 2021), improving performance over comparable baselines. Finally we show that the model can be adapted to generalize to new unseen entities by updating the memory, without retraining.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Our method represents knowledge in a corpus as a collection of "mention encodings" -dense vector representations for every entity mention that appears in the corpus. Every time an entity appears in a passage -"[Barack Obama] was elected president in 2008" -some property of the entity or its relation to other entities is described. The first component of our method, the Mention Encoder model, is responsible for distilling information from entity mentions in the corpus into high-dimensional mention encodings. We use the Mention Encoder to encode each entity mention in English Wikipedia and gather encodings into a Mention Memory. The purpose of the Mention Memory is to capture all knowledge contained in the corpus in a way that can be easily integrated into a Transformer. The second component of our method, the TOME model, applies sparse attention over the Mention Memory to incorporate external information from the corpus into a Transformer model. An overview of the whole method is shown in Figure 1.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Jointly training the Mention Encoder and TOME models is computationally costly, since it would require backpropagating through the Mention Encoder for each attended mention. Consequently, we propose to train the models in two stages. First, we pre-train the Mention Encoder and generate the Mention Memory. Second, we pre-train the TOME model while keeping the Mention Memory frozen: the gradient does not propagate through it and the memories are not modified. Mention Encoder pre-training is specifically designed such that mention encodings capture relevant contextual information about each mention and are useful for TOME even without joint training. We formally define these models in sections 2.1 and 2.2, and their pre-training procedures in 2.3 and 2.4.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Notation. An input to the model is a passage x = x 1 , . . . , x T of length T . We assume that each passage has been annotated with an NER system. Following Baldini Soares et al. (2019) we use special entity markers to highlight entity mentions in the passage. We introduce tokens [E start ] and [E end ] to the vocabulary and insert them before and after each mention in the passage. For example, the original passage "What is the nationality of the hero who killed Medusa" turns into "What is the
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Each mention m in a passage is described by a tuple (s, e), where s and e are start and end positions of the mention. We consider entity markers to be part of the corresponding mention, so that x s = [E start ] and x e = [E end ]. Representations of these tokens are later used to generate mention encodings. (1) The Mention Encoder is a Transformer model with two final SpanEncodingLayers that produce key and value mention encodings. Value mention encodings store context-level information about each mention and are used as inputs to the TOME model. Key mention encodings identify the type of information stored in the value encodings and serve as attention keys for the memory layer. These two SpanEncodingLayers do not share weights.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
After the Mention Encoder is pre-trained (see section 2.3), we use it to generate a Mention Memory from entity mentions in Wikipedia. While we could include encodings of any corpus mention in the Mention Memory, we focus on grounded mentions which can be linked to Wikipedia entities. We denote these as linked mentions, which we hypothesize contain information that can be retrieved and grounded. We gather mention encodings into matrices MemKey ∈ R N ×d K and MemValue ∈ R N ×d V , where N is the total number of linked entity mentions in English Wikipedia (approximately 150 million) and d K and d V are dimensions of key and value encodings. Additionally, we record entity (Wikipedia) IDs of mentions in MemEnt ∈ R N , which we use as labels for auxiliary losses, not as inputs to the model or supervision on retrieval. MemKey(i), MemValue(i), MemEnt(i) correspond to the key encoding, value encoding and entity ID for the i-th linked mention in Wikipedia.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
The TOME model incorporates information from a text corpus into a Transformer by applying sparse attention over the Mention Memory. The model consists of one or more TOMEBlocks, each containing a memory attention layer followed by a post-processing Transformer block. Memory attention layers retrieve and attend to relevant "memories" for every mention in the input passage. The model then processes the retrieval-augmented representation with the Transformer block, allowing it to access and combine information from multiple sources in the corpus. Finally, multiple TOMEBlocks enable the model to refine retrievals and perform multi-hop reasoning. More formally, a TOMEBlock receives the output representation of the previous layer H and produces new representations H M = MemoryAttention(H),
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
The TOME model encodes input passages x with the word embedding layer and initial Transformer block and then applies one or more TOMEBlocks
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
In this work we consider two configurations of the TOME model: TOME-1 and TOME-2, with one and two TOMEBlocks respectively. Each TOMEBlock of TOME-2 contains half as many Transformer layers as in TOME-1 to hold the total number of Transformer layers fixed between models.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Each memory attention layer is implemented as a sparse dot-product attention layer that takes the output H of the previous Transformer block, incorporates information from the Mention Memory, and returns a representation M (omitting layer indices). Consider a mention m that starts at position s and ends at position e. We start by computing its query mention encoding Query(m) by applying a SpanEncodingLayer Query(m) = SpanEncodingLayer(H, (s, e)),
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Query mention encodings are used to retrieve relevant memories from the Mention Memory table.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
However, applying standard attention over 150 million mention encodings is infeasible. Instead, we first perform approximate nearest neighbor search to retrieve the top-K mentions with the largest dot product between query Query(m) and key mention encoding from MemKey. We denote the set of these memories as TopMem(Query(m)). We compute attention over these memories and incorporate the result into the token contextual representation at position s
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Value(m) = i∈TopMem(Query(m))
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
where W U is a learnable matrix of shape d × d V .
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Approximate nearest neighbor search (ANNS) can be performed cheaply using one of multiple ANNS libraries, for example ScaNN (Guo et al., 2020). We implemented two on-device search methods to avoid the engineering complexity of real-time communication with an ANNS server, though we have verified this is also viable. The first naively computes a simple dot-product between passage queries and memory keys, and was used in our main experiments as it was easiest to implement. We also implemented and will be releasing a much faster version based on CPU ANNS methods. The memory is sharded over devices, so that the device-memory overhead is negligible.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Holding the number of entries in memory fixed, the compute cost of retrieval from memory does not grow with the size of the reader or the dimensionality of the memory values, so that the relative cost of the memory layer becomes smaller with reader size. In particular, the overhead from the memory used in our pre-training setting is small for BERT-Large and up. More details on ANNS implementation and overhead can be found in Appendix C.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
While backpropagating through a Wikipedia-scale mention memory is challenging, it is possible to train smaller-scale memory architectures end-to-end. We take an approach inspired by MARGE (Lewis et al., 2020a) and READTWICE (Zemlyanskiy et al., 2021) which apply crossattention over documents within a batch. In particular, we process passages in each batch twice. As a first step, the Mention Encoder model generates mention encodings from each passage and aggregates the mention encodings into a batch-wide memory table. In the second step, we apply a TOME architecture that attends to the batch memory, which we call BATCH-TOME. Note that BATCH-TOME is just used for pre-training the Mention Encoder and not evaluated on any downstream tasks.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Mention Encoder and BATCH-TOME are jointly trained end-to-end so that the Mention Encoder is encouraged to produce mention encodings that contain useful information for BATCH-TOME.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
We want to make sure the batch memory contains relevant mentions, so we pre-train the models on batches of passages constructed from related Wikipedia articles with high entity overlap. Appendix A.1 provides more details on Mention Encoder data generation. We use the pre-trained Mention Encoder to construct the Mention Memory table from corpus, and use the BATCH-TOME model as the initialization point for TOME-specific pre-training (described in Section 2.4).
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Masked language model. Our primary pre-training objective is the standard masked language modeling task, with the loss computed based on the output of the second read (BATCH-TOME). To encourage the model to rely on memory, we increase the task's difficulty relative to standard BERT pre-training by masking entity mention tokens more aggressively.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Coreference resolution. We wish to encourage the Mention Encoder to represent the entity attributes expressed by entity mentions, so we also employ an entity-oriented pre-training task to the output of BATCH-TOME for which such attribute information is likely to be especially helpful. Unlike Entities as Experts (Févry et al., 2020), BATCH-TOME does not use entity embeddings, so we cannot use the entity linking task. Instead, we apply a related entity coreference resolution objective, which asks the model to predict whether two linked mentions correspond to the same entity based on the similarity of their encodings. Given that entity surface forms are frequently masked, the model needs to instead use the properties of other mentions in the batch to determine which entity it is most compatible with, incentivizing the Mention Encoder to encode such properties. We compute a coreference mention encoding for every linked mention in the batch by applying a separate SpanEncodingLayer on the output of BATCH-TOME. The loss is implemented using cross-entropy over dot-product similarity scores. See Appendix A.2 for details.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
As TOME attends to the full Mention Memory instead of in-batch memory, we do not employ the batching procedure from Mention Encoder pre-training, instead sampling Wikipedia passages randomly. For the same reason, we replace the in-batch entity coreference objective by Mention Memory entity coreference, in which the model has to predict which mentions from the Mention Memory share an entity with the input mention. The goal of this auxiliary objective is to incentivize the model to learn to retrieve informative mention encodings to solve the semantically challenging task. Mention Memory entity coreference also allows us to solve tasks like TriviaQA or ComplexWebQA without a decoder by directly predicting the answer entity.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Entity prediction. Analogous to batch coreference resolution loss we compute mention encoding z m using the output of the TOME model. As in section 2.2, TopMem(z m ) returns the top K memories with the largest dot product between the mention encodings z m and key mention encodings MemKey from the Mention Memory. The score EntProb(m, j) of entity j equals the sum of attention weights of memories corresponding to this entity.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
The final entity prediction is arg max j EntProb(m, j). Entity prediction loss L ep (m) for a mention m of entity Ent(m) is L ep (m) = − log EntProb(m, Ent(m)). Total loss equals the average loss over linked input mentions for which at least one memory of the same entity is retrieved.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Disallowed same passage retrieval. For each passage in the pre-training corpus, there exist memories corresponding to mentions in the passage generated from the unmasked version of the same passage. In order to prevent the model from 'cheating' by attending to such memories, we set the attention weight for all memories from the same passage to zero.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Our approach lies at the intersection of three lines of work: i) knowledge-augmented language models, ii) employing a text corpus as a virtual knowledge base, iii) retrieve-and-read methods.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Knowledge-augmented language models. Entities as Experts (EaE) (Févry et al., 2020) injects information into a Transformer model model with an intermediate attention layer over trainable entity embeddings, which serve as an aggregate representation of entity information in a text corpus. In contrast, TOME attends to a much larger table of mention encodings, allowing for retrieval of more fine-grained information. Attending to mentions as opposed to entity representations also enables TOME to generalize to unseen entities. FiLM  extends EaE by adding an attention layer over facts from a KB on the output of the Transformer. The fact attention layer enables more fine-grained queries but still retrieves aggregate entity embeddings as values, which are also not reasoned over by the Transformer. KnowBERT (Peters et al., 2019) is similar to EaE, but with entity embeddings generated from a KB instead of trained end-to-end with a text corpus. MARGE (Lewis et al., 2020a) and READTWICE (Zemlyanskiy et al., 2021) incorporate dense representations from other passages within the same batch into a Transformer through sparse top-k attention. The first pre-training stage of our method for training the Mention Encoder is similar to MARGE and READ-TWICE. However, TOME performs global attention over a full corpus, rather than a single batch. Furthermore, TOME attends to a Mention Memory consisting of pre-computed dense representations. Therefore TOME is not limited to downstream task with batches of relevant documents, and does not need to apply an expensive reader model to an entire batch of documents for each input.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Text corpus as virtual knowledge base. DrKIT (Dhingra et al., 2020) performs multi-hop question answering by using a text corpus as a virtual knowledge base. Similar to TOME, the authors apply a mention encoder to convert the corpus into a table of mention encodings. A Transformer model encodes the question into dense queries, which are compared with the mention encodings to traverse the VKB. Conversely, TOME retrieves mention encodings, and then jointly processes them inside the Transformer. In follow-up work to DrKIT, OPQL ) uses a FiLM-like approach to access a memory of relation mentions, which are encoded with a self-supervised relation encoder. However, the relation mention encoding combines a mention-specific relation representation with EaE-like entity encodings, so they are less fine-grained than TOME's encodings. Unlike TOME, OPQL also lacks a sparse large-scale retrieval mechanism, and relies on ad hoc heuristics to limit the size of the memory. 1 MOLEMAN (FitzGerald et al., 2021) compares a passage mention encoding with mention encodings from a corpus to perform entity linking, but does not retrieve the mentions.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Retrieve-and-read methods. REALM (Guu et al., 2020) learns to retrieve relevant passages from a text corpus in a self-supervised manner. Retrieved passages are concatenated to the input passage which is then re-encoded by a Transformer model to perform a task. The key difference between retrieve-and-read approaches (Guu et al., 2020;Karpukhin et al., 2020;Lewis et al., 2020b;Izacard & Grave, 2021) and TOME is that TOME retrieves dense representations, as opposed to text. That means that TOME only applies a reader model once to a single input, while retrieve-and-read approaches have to apply an expensive BERT reader to many different passages. In addition, Transformer models can only process relatively short sequences, which imposes a binding constraint on the number of retrieved text passages that can be processed together, whereas TOME can retrieve and reason over information from many sources inside the same reader. Generative models like RAG (Lewis et al., 2020b) or FiD (Izacard & Grave, 2021) attend to different retrieved documents in the decoder, but still have to apply a BERT read for every retrieved document, do not consider interaction between retrievals while encoding the question, and cannot perform iterative retrieval.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
The Mention Encoder is based on a BERT-base model with two final SpanEncodingLayers that produce key and value encodings. Mention Encoder and BATCH-TOME share Transformer weights during Mention Encoder pre-training. The Mention Memory consists of mention encodings for N = 150 million linked Wikipedia entity mentions. Transformer layers in TOME and BATCH-TOME models are equivalent to those in the BERT-base model. The TOME InitialTransformerBlock contains 4 Transformer layers. TOME-1 has a single TOMEBlock with 8 Transformer layers, and TOME-2 has two TOMEBlocks with 4 Transformer layers each. Therefore, the number of trainable parameters in TOME-1 and TOME-2 is approximately the same as in BERT-base. We use a smaller Mention Memory containing 38m uniformly sampled memories for TOME pre-training. During fine-tuning and evaluation we utilize the full Mention Memory. Appendix A contains more details.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
We compare TOME with existing methods that utilize textual information from a corpus in a language model. These can be divided into generative LLMs (T5), entity embedding retrieval (Entities as Experts, OPQL), extractive retrieve-and-read (REALM) and generative retrieve-and-read (RAG, Fusion-in-Decoder). TOME occupies a novel position in the space of retrieval models, being more fine-grained than entity embedding retrieval methods, but performing all its reasoning with a single BERT read, unlike retrieve-and-read methods. The most closely comparable models are Entities as Experts and REALM, and we use these as our primary baselines. We report other baselines for reference, with the caveat that these results are not apples-to-apples: RAG and Fusion-in-Decoder have large decoders and retrievers and encode a large number of passages with a BERT reader for each question compared to TOME's single read. Fusion-in-Decoder and RAG 2 also use ground-truth supervision for retrieval. We mark the number of parameters and BERT applications for each baseline in the result tables. Consistent with retrieve-and-read, we count the parameters of the Mention Encoder and TOME, but not the size of the non-trainable and sparsely accessed Mention Memory.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Data. Our first set of experiments evaluates TOME on the claim verification tasks FEVER (Thorne et al., 2018), HoVer (Jiang et al., 2020), and FM2 (Eisenschlos et al., 2021) in which the model is provided with a claim and has to determine whether the claim is supported by the Wikipedia corpus. FEVER is a larger dataset with 186k claims for which most of the claims can be verified with a single Wikipedia passage. In contrast, HoVer is smaller with 26k claims, but is explicitly constructed to require evidence from multiple sources and multiple reasoning steps. FM2 is also smaller and is constructed through an adversarial game that leads to more challenging retrieval. The claim verification training data contains gold evidence passages, but unlike most published results we do not use these, leaving only the accuracy of the claim verification to guide the retrieval.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Results. Table 1 contains our claim verification results. TOME outperforms both Entities as Experts and REALM, especially on HoVer and FM2. This is consistent with the properties of TOME: HoVer requires combining detailed information from multiple sources, which TOME is especially well equipped to do compared to aggregate entity-based or retrieve-and-read models. FM2 features generally challenging retrieval and may benefit from contextualizing retrieved evidence.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Data. In a second set of experiments we evaluate TOME on TriviaQA (TQA) (Joshi et al., 2017), ComplexWebQuestions (CWQ) (Talmor & Berant, 2018) and EntityQuestions (EQ) (Sciavolino et al., 2021), open-domain QA tasks for which most answers are Wikipedia entities. We approach these datasets as entity-linking tasks, as in Févry et al. (2020). We append a mask token to each question, which is marked as a question mention. The probability for each candidate entity is predicted as the aggregate attention weight on mentions of the entity (Section 2.4). Questions with answers that do not correspond to entities in our entity vocabulary are marked as answered incorrectly. TQA  Results. Table 2 contains the results for TQA, CWQ and EQ experiments. Like TOME, Entities as Experts and OPQL treat the above datasets as entity-linking tasks. REALM performs extractive QA, while T5, RAG and Fusion-in-Decoder generate the answer. We note a similar pattern of results as for claim verification. TOME strongly outperforms Entities as Experts on all tasks. TOME performs slightly better than REALM on a simple task like TriviaQA (entity subset) and strongly outperforms REALM on more challenging tasks that require multiple (CWQ) or challenging (EQ) retrieval.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
What memories does TOME retrieve? Given that TOME retrieval is unsupervised, it is natural to ask what memories it learns to retrieve. First, we observe that BATCH-TOME and TOME trained on just the MLM objective learn to attend to memories of the same entity as the passage linked mention (55% and 41% average attention score). This is promising as entity mentions from the same entity often contain mutually relevant information. Quantitative evaluation of downstream retrieval is challenging as TOME often retrieves mentions that are not part of, but equally informative as gold passages. Instead, we provide TOME retrievals on the first three samples of the HoVer dev set to demonstrate its retrieval behavior without cherry-picking. Table 3 demonstrates a successful simple  retrieval, while Table 4 displays interesting multi-hop retrieval. The last is found in Appendix D.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Importance of memory size. Figure 2 shows claim verification performance as a function of memory-size during fine-tuning (pre-training memory size is held constant). For smaller memory sizes, entries in memory are uniformly sampled from the full Mention Memory. Performance increases smoothly with memory size. Larger memory size yields diminishing returns, perhaps reflecting that entity mentions may contain overlapping information. Table 4: TOME-2 retrievals for the first HoVer dev sample. We show top-1 retrieval results for the first (−→ 1 ) and the second (−→ 2 ) memory attention layers for passage mentions "Life Goes On" and "Hungry" 3 . Memory mentions are in brackets. The first retrieval for the "Life Goes On" is a different song with the same name and the first retrieval for "Hungry" is related but not useful. However, the second retrieval for "Life Goes On" identifies the correct song and describes its position on the album while the second retrieval for "Hungry" captures its position relative to "Life Goes On".
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Claim: The song recorded by Fergie that was produced by Polow Da Don and was followed by Life Goes On was Hungry. Label: TRUE Life Goes On −→ 1 . . . and Johnny J produced the chart topping hits "All Bout U", "How Do U  Table 5: Accuracy on held-out subset of Trivi-aQA and ComplexWebQuestions (CWQ) questions. TOME-1-unseen was pre-trained and finetuned with memory without entities from heldout set and evaluated with full memory. Note that performance is considerably lower than on the full dev set as answers in the held-out set (which are in dev but not train) are more likely to be rare entities.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
TriviaQA dev CWQ dev TOME-1 17.4 16.4 TOME-1-unseen 17.6 16.7
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Zero-shot transfer to unseen entities. An important advantage of memory architectures is that the behavior of the model can be steered by deciding what to include in the memory. Here we show that the TOME model can use information that was not present in memory during training. We sample questions in the TQA and CQA dev sets, and generate a subset of the memory without any mentions corresponding to the answer entities for those questions. Then we pre-train and fine-tune a model on this smaller memory, which we call TOME-unseen. We evaluate TOME-unseen on the sampled questions using the full memory for evaluation only, and compare to standard TOME. Table 5 shows that using full memory only during evaluation does not lower performance.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
We introduced TOME, a Transformer model that performs attention over a semi-parametric representation of the entire Wikipedia text corpus. This representation, or Mention Memory, consists of a dense encoding for each entity mention in Wikipedia. TOME can retrieve information from multiple sources without supervision, aggregate information within the Transformer, and reason over the retrieved information. TOME leads to strong improvements on multiple open-domain claim verification and entity-based question answering tasks.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
We train on English Wikipedia, processed with the entity linking and named entity recognition tools from the Google Cloud NLP API 4 . We use existing hyperlinks in Wikipedia as additional entity annotations. All models are pre-trained on 128 TPUs using AdamW optimizer (Loshchilov & Hutter, 2019) with learning rate 1e-4 and batch size of 4096. Each passage in the batch has length T = 128, excluding entity tokens. The Mention Encoder and BATCH-TOME are pre-trained for 1 million steps with 50k warmup steps, and TOME is trained for 500k additional steps with 25k warmup steps after initialization from BATCH-TOME. Both models are trained with linear learning rate decay. Mention Encoder and BATCH-TOME share Transformer weights during Mention Encoder pre-training. We apply gradient clipping with a norm of 1.0 and weight decay of 0.01. Weight decay is applied to all weights except layer norm and bias weights.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
BATCH-TOME and TOME are trained with weight 0.85 on the MLM objective and 0.15 on the entity coreference resolution objective. We mask 20% of whole entity mentions and 10% of other tokens. We limit the coreference resolution objective to mentions of the 1 million most frequent Wikipedia entities. We use 24 mentions per sample, with a batch size of 32 samples per TPU. We subsample mentions uniformly if the average number of annotated mentions on a TPU exceeds 24. Key mention encoding have dimension d K = 128 and value and coreference mention encodings have dimension d V = d C = 512.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Disallowed same passage retrieval for Mention Encoder. We want the model to use memory as a source of additional information for processing a passage. Therefore, we explicitly set attention weights to 0 for memories generated from the same passage as the current one.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
We pre-train Mention Encoder to produce mention encodings that are useful for BATCH-TOME. In order to provide BATCH-TOME with an incentive to use the memory, we need to ensure that mentions from different samples within a batch are relevant to each other. We achieve this by batching passages from the same or related Wikipedia articles.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
We generate clusters of 256 passages from Wikipedia articles using a greedy method. First, we create a cluster from the longest unused Wikipedia article and add related articles until the cluster consists of 256 passages. In particular, at each step we add the article with the largest Jaccard similarity between its entity set and the entity set of articles in the current cluster.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
For every linked mention m in the batch we compute a mention encoding z m by applying a separate SpanEncodingLayer on the output of the BATCH-TOME. First, we compute the loss for every linked mention m in the batch. To this end, we denote linked mentions in every other passage in the batch as positive, P + (m), if they have the same entity ID as m and negative, P − (m), otherwise. The loss per mention is an average of cross-entropy losses per every positive mention m + ∈ P + (m)
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
The total loss is average of losses per linked mentions which have at least one positive mention (set P + (m) is not empty).
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
B.1 FINE-TUNING SETUP TOME is fine-tuned on 32 TPUs using the Adam optimizer with a learning rate of 1e-5 and total batch size 32. In contrast to pre-training we set max mentions to 32 per sample for fine-tuning. We use 1000 warmup steps and linear learning rate decay. Gradient clipping and weight decay are the same as during pre-training. We take the highest scoring checkpoint on dev sets and evaluate it on the test set. We use the spaCy noun chunker to detect noun phrases and treat these as claim/question entity mentions.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
The model can be fine-tuned with full memory on a server of 8 A100 GPUs or 16 v3/v4 TPUs. A model with half memory (75M mentions) can be fine-tuned on 8 V100s/P100s or 8 TPUs.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Following Guu et al. (2020) we used REALM to perform extractive question answering on TriviaQA and ComplexWebQuestions datasets. We also adapted the model to the classification setting in order to apply it to claim verification tasks. Given an input claim X we compute probability of a prediction Y (whether the claim holds true or not) as a marginalization over retrieval Z.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
where Pr(Y |X, Z = z) is the output probability produced by the reader model and Pr(Z = z|X) is produced by the retrieval model.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
See Table 6 for the results on development and test splits of claim verification datasets. Additionally, Table 7 compares our FM2 results to the original dataset baselines. We report additional results on the EntityQuestions dataset from Sciavolino et al. (2021). The dataset consists of questions involving rare entities, making it especially challenging for modern retrievals methods such as DPR. Evaluation results for TOME models and baselines are shown in Table 8 and Table 9. Following (Sciavolino et al., 2021) we report recall at 20 as an evaluation metric. Since TOME retrieves mentions rather than passages, a direct comparison is difficult. We evaluate TOME conservatively, treating recall at 20 as successful if one of the 20 highest scoring mentions belongs to the correct entity (in contrast to DPR, for which the correct answer only has to be somewhere in the retrieved 100 word document). TOME sets the state of the art on this dataset and outperforms DPR by a very large margin. REALM cannot be fairly compared to DPR due to longer retrieved passages (100 vs 288 tokens). Therefore, we perform a separate experiment using accuracy with REALM as a baseline, showing large performance gains over REALM as well.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
We perform several ablation experiments for the pre-training procedure (see Table 10). First, results show that the entity prediction objective (c.f. Section 2.4) is not essential for TOME pre-training. Performance on claim verification datasets (FEVER and HoVer) is not affected by whether we use entity prediction for pre-training. More surprisingly, removing this objective only slightly decreases the performance on entity question answering datasets (TriviaQA and ComplexWebQuestions). We predict entities for question-answering in the same way as we do for the entity prediction objective during pre-training (c.f. Equation 10), so we expected the entity prediction auxiliary loss to be important.   On the other hand, a related entity coreference objective (c.f. Section A.1 and Appendix A.2) is crucial for Batch-TOME and Mention Encoder pre-training. That is consistent with our intuition that semantically challenging tasks incentivize the model to store useful information in memory.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
We initialize TOME model with a pre-trained BATCH-TOME model which we find to be especially important for warming up retrieval. If TOME is initialized from scratch (or even from BERT weights), TOME does not learn to use the memory. In fact, TOME has to be initialized from the same BATCH-TOME used to generate the memory. This implies that multi-stage training is a vital ingredient for TOME to succeed. Our explanation for why TOME is sensitive to initialization is that TOME needs to learn two skills: first, to effectively use retrieved mentions for its predictions, and second, to retrieve relevant mentions. Learning both capabilities end to end gives rise to a mutual dependence: to get a signal for learning how to use retrieved mentions, the retrieved mentions have to be useful, and to learn to retrieve useful mentions, the model needs to utilize retrieved mentions. If initialized from scratch, the model is not able to learn both skills simultaneously. The pre-training stage with the smaller in-batch memory functions as a curriculum to address that problem.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Nearest neighbor search is an extremely common problem, and there exist numerous approaches and packages for fast approximate nearest neighbor search (ANNS). (Guo et al., 2020;Johnson et al., 2017). Most approaches employ two methods for fast search: 1) compress the search table through projecting to a lower dimension and quantization and perform comparisons in this compressed space and 2) divide the search table in buckets of similar items, and search only a subset of the buckets. Retrieve-and-read use ANNS packages to search for related passages (Guu et al., 2020;Lewis et al., 2020b).
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Applying such packages for TOME is slightly trickier, as TOME needs to perform ANNS inside the model. One viable route is to compute queries on-device, transmit them to a separate ANNS server and them transmit them back. We would recommend this approach for GPU accelerators, with faster host-device communication and slower device-device communication. As we are using TPU accelerators, we decided to use on-device ANNS, which does not require coordinating additional servers and will potentially allow for backpropagating through memory in future work.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
We shard the Mention Memory over all TPU devices. We perform search by distributing each query to all devices and retrieving top-K search results from each local memory shard. Then, the results are distributed back to the original devices and the local search results are aggregated through another, global top-K.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
The first method we describe is naive dot-product search, taking advantage of matrix multiplication capacity of TPU accelerators. In this method we perform search over local shards by taking the dot product between the query and the local memory shard and performing an approximate top-k operation over the results. Dot-product search is easy to implement and fast for smaller memory sizes (up to 10 million entries). We implemented this method first due to its simplicity and our primary experimental results employ this search method.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
To speed up search we implemented method 2) from standard CPU-based ANNS, bucketing the search table and searching only a subset of buckets. In particular we perform k-means clustering to divide the Mention Memory into clusters, and perform dot-product search over the top n s clusters on each device.
Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
While the Mention Memory is stored on-device, memory overhead is negligible as the memory table is sharded. For pre-training the Mention Memory took up 2.2% of available device memory. Table  11 shows percentage of time spent on ANNS in TOME-1 pretraining for different reader architectures. The relative overhead of search becomes smaller with reader size, and ANNS overhead in particular becomes negligible for BERT-Large and up. We did not measure standard CPU ANNS overhead, but it should be comparable to or faster than our ANNS numbers. For ANNS in TOMEBlocks we take top-2 search results from each local memory shard, and apply top-128 over the retrieved results. For ANNS in the entity prediction layer we take top-32 search results from each local shard, and aggregate across shards without applying an additional top-K operation. Table 12: TOME-2 retrievals for the second HoVer dev sample. We show top-1 retrieval results for the first (−→ 1 ) memory attention layer for passage mentions "the novel", "the movie" and "the album". Memory mentions are in brackets. We can see that the model can retrieve relevant mentions for non-named passage mentions, and generally understands it is looking for mentions related to music. However, while the best retrieval for "album" is from a passage that mentions sampling the Shining, it is quite far removed and it is likely the retrieval is not sufficiently accurate here. 
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
Pre-trained language representation models (PLMs) such as ELMo (Peters et al., 2018a), BERT (Devlin et al., 2019a) and XLNet  learn effective language representations from large-scale nonstructural and unlabelled Work in progress. corpora and achieve great performance on various NLP tasks. However, they are typically lack of factual world knowledge (Petroni et al., 2019;.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
Recent works (Zhang et al., 2019;Peters et al., 2017;Liu et al., 2019a) utilize entity embeddings of large-scale knowledge bases to provide external knowledge for PLMs and improve their performance on various NLP tasks. However, they have some issues: (1) They use fixed entity embeddings learned by a separate knowledge embedding (KE) algorithm, which cannot be easily aligned with the language representations because they are essentially in two different vector spaces. (2) They require an entity linker to link the words in context to corresponding entities so that they can benefit from the entity embeddings, which makes them suffer from the error propagation problem. (3) Their sophisticated mechanisms to retrieve and use entity embeddings lead to additional inference overhead compared with vanilla PLMs.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
Actually, knowledge embedding methods have a strong connection with NLP models. There are not only many works integrating knowledge embeddings into NLP models to improve the performance of NLP applications such as machine translation (Zaremoodi et al., 2018), reading comprehension (Mihaylov and Frank, 2018;Zhong et al., 2019) and dialogue system (Madotto et al., 2018), but also some early works use text as additional information (Xie et al., 2016;An et al., 2018) or jointly train the knowledge and text embedding in the same space (Wang et al., 2014;Toutanova et al., 2015;Han et al., 2016;Cao et al., 2017Cao et al., , 2018.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
In this paper, we propose to learn knowledge embedding and language representation with a unified model and encode them into the same semantic space, which can not only better integrate knowledge into PLMs but also help to learn more informative knowledge embeddings with the effec-tive language representations. We propose KE-PLER, which is short for "a unified model for Knowledge Embedding and Pre-trained LanguagE Representation". We collect informative textual descriptions for entities in the knowledge graph and utilizes a typical PLM to encode the descriptions as text embeddings, then we treat the description embeddings as entity embeddings and optimize a KE objective function on top of them. The key idea is to encode structural knowledge in the textual representation of entities using a PLM, which can generalize to unobserved entities in the knowledge graph.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
Our KEPLER enjoys the following advantages:
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
(1) We integrate world knowledge into PLMs with the supervision of the KE objective, which is more flexible for the PLMs, and encode the entity and text into the same space, which avoids the gap between the language representations and fixed entity embeddings.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
(2) We do not need an entity linker or additional mechanisms to retrieve corresponding entity embeddings, which avoids the error propagation problem and extra overhead. During inference, our KEPLER is exactly the same as standard PLMs, which can be adopted in a wide range of NLP applications. (3) Different from conventional KE methods, our KEPLER encodes textual entity descriptions as entity embeddings, which enables our model to infer knowledge embedding in the inductive setting (get entity embeddings for the unseen entities). This is especially useful for deployment, where the model may deal with unseen entities.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
The existing KE datasets are relatively smallscale, which is not sufficient to pre-train a large model, and typically lack of description data and a data split for the inductive setting. Therefore, we construct Wikidata5m, a new large-scale knowledge graph dataset with aligned text description for each entity. Wikidata5m is a subset of Wikidata (Vrandečić and Krötzsch, 2014), a free knowledge base with about sixty million entities. To ensure each entity is informative and the knowledge base is as clean as possible, we only select the entities with corresponding Wikipedia pages. Our Wikidata5m contains five million entities and twenty million triplets. We also benchmark several classical KE methods on Wikidata5m to facilitate future research. To our knowledge, this is the first million-scale general knowledge graph dataset.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
To summarize, our contribution is three-fold:
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
(1) We propose to encode entities and texts into the same space and jointly train the KE and language modeling objectives, and then get a better knowledge-enhanced PLM which avoids error propagation and additional overhead. Experimental results on various NLP tasks demonstrate the effectiveness of our KEPLER. (2) We encode textual descriptions as entity embeddings, which improves KE with textual information and enables inductive KE. (3) We introduce a new large-scale knowledge graph dataset Wikidata5m, which may promote the research on large-scale knowledge graph, inductive knowledge embedding and interactions between knowledge graph and NLP.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
… is an independent agency … for the civilian space program … Kepler space telescope … is a retired space telescope launched by NASA to … Named after astronomer Johannes Kepler. Knowledge Graph Embeddings In recent years knowledge embeddings have been extensively studied through predicting missing links in graphs. Conventional models define score functions for relation triples (h, r, t) and predict head or tail entities with scores of candidate entities. For example, TransE (Bordes et al., 2013) treats tail entities as translations of heads, while DistMult (Yang et al., 2015) use matrix multiplications as score functions and ComplEx (Trouillon et al., 2016) adopt complex operations based on it. RotatE (Sun et al., 2019a) combines the advantages of both of them.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
Among these works, Xie et al. (2016) propose to utilize entity descriptions as an external information source and introduce an entity description encoder to enhance the TransE score function. Though similar to our method, Xie et al. (2016) aim at utilizing entity descriptions to help knowledge representation learning, while we take entity descriptions as a tool to incorporate external knowledge in our model.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
In this section, we introduce the structure of our KEPLER model, and how we combine two training goals of masked language modeling and knowledge representation learning.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
To incorporate world knowledge into our pretrained language representation models (PLMs), we design a multi-task loss as shown in Figure 1 and Equation 1,
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
where L KE represents knowledge embedding loss and L LM represents language model loss. Since our PLMs are involved in both tasks, jointly optimizing the two objectives could implicitly integrate knowledge from external graphs with text encoders, while keeping the strong abilities of PLMs for syntactic and semantic understanding.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
More specifically, we adopt a general L KE format using negative sampling,
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
where (h, r, t) is the correct triple from knowledge graphs and (h i , r, t i ) are negative sampling triples. d r is the score function, for which we have many choices. Different from conventional knowledge embedding methods, for entity embeddings h and t, instead of looking up in embedding tables, we use PLMs as our text encoders to extract entity representations from their descriptions.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
For L LM , many alternatives for pre-trained language representation can be used, e.g., masked language model (Devlin et al., 2019b). Note that those two tasks only share the text encoder and for each mini-batch, text sampled for L KE and L LM is not (necessarily) the same.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
Though we have many alternatives of model structures and training objectives to choose under KE-PLER framework, here for better clarity, we introduce a specific one that we use in experiments.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
We use the transformer architecture (Vaswani et al., 2017) as in (Devlin et al., 2019b;Liu et al., 2019c), which we will not address in details. To be more specific, we use RoBERTaBASE codes and checkpoints 1 in all our experiments since it is one of the state-of-the-art pre-trained models with acceptable computing requirements. Besides the training data and hyperparameters, one of the major differences between RoBERTa and BERT is that RoBERTa uses Byte-Pair Encoding (BPE) (Sennrich et al., 2016) to better tokenize rare words.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
Given a sequence of tokens x 1 , x 2 , ..., x N , the input format is PLM Objective Inspired by BERT (Devlin et al., 2019b), MLM randomly selects 15% of input tokens, among which 80% are masked with the special mark [MASK], 10% are replaced by another 1 https://github.com/pytorch/fairseq random token, and the rest remain unchanged. Under MLM, models try to predict the correct tokens and a cross-entropy loss is calculated over the selected positions.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
We adopt the pre-trained checkpoint of RoBERTaBASE for the initialization of our model. However, we still keep MLM as one of our objectives to avoid catastrophic forgetting (McCloskey and Cohen, 1989) while training towards the KRL loss. Note that experiments show that only further pre-training from RoBERTaBASE checkpoint does not bring promotion, suggesting that the combination of the two tasks contributes most to the performance.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
We use the loss formula from (Sun et al., 2019b) as our KE objective, which takes negative sampling (Mikolov et al., 2013) for efficient optimization:
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
where (h, r, t) is the correct triple, (h i , r, t i ) are negative sampling triples, γ is the margin, σ is the sigmoid function, and d r is the score function, for which we choose to follow TransE (Bordes et al., 2013) for its simplicity and efficiency,
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
where we take the norm p as 1. Due to the limit of computing resources, we take the negative sampling size n as 1. The negative sampling policy is to fix the head entity and randomly sample a tail entity, and vice versa.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
Different from conventional KE methods, we do not have an entity embedding lookup table. Instead, we use our KEPLER model to encode the corresponding entity descriptions and take the [CLS] outputs as the entity embeddings.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
Like all BERT-like models, we fine-tune KEPLER on downstream tasks and use [CLS] output for sentence-level prediction and the outputs of all tokens for sequence labelling tasks (Devlin et al., 2019b). For supervised relation extraction and fewshot relation extraction, we follow the approaches from (Baldini Soares et al., 2019) and (Gao et al., 2019) respectively. 
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
We construct a new large-scale knowledge graph dataset with aligned text descriptions. Our dataset is built by integrating Wikidata (Vrandečić and Krötzsch, 2014), a large-scale open knowledge base, with Wikipedia. Each entity in the knowledge graph is aligned with its text description in Wikipedia pages. In the following sections, we will first introduce the data collection steps, and then give the benchmarks of popular KE methods on this dataset.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
We pull the latest dump of Wikidata 2 and Wikipedia 3 from their websites respectively. We remove pages whose first paragraphs contain fewer than 5 words. For each entity, we align it to a Wikipedia page with the MediaWiki wbgetentities action API. The first section of Wikipedia pages is extracted as the description for entities. Entities that have no corresponding Wikipedia pages are discarded.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
To construct the knowledge graph, we retrieve all the statements in entity pages, and map the entities and relations in statements to their canonical IDs in Wikidata. A statement is considered to be a valid triplet if both of its entities can be aligned with Wikipedia pages and its relation has a non-empty page in Wikidata. The final knowledge graph dataset contains 4,813,455 entities, 822 relations and 21,344,269 triplets, where each entity has a text description. Statistics of our Wikidata5m dataset and four widely-used datasets are showed in Table 1. Top-5 entity categories are listed in Table 3. We can see that our Wikidata5m is much larger than existing knowledge graph datasets, covering all sorts of domains.   
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
The data split statistics for the conventional transductive setting are also shown in Table 1.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
In this work, we also evaluate models on the challenging inductive setting, which requires the models to produce entity embeddings for entities which are not seen at the training time and also do link predictions for the unseen entities. So we provide a data split for the inductive setting evaluation. The statistics for the inductive setting data split are shown in Table 2. In the inductive setting, the entities and triplets in training, validation and test sets are mutually disjoint, while in the transductive setting, only the triplet sets are mutually disjoint.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
To assess the challenges of Wikidata5m, we benchmark several popular knowledge graph embedding models on the dataset. Since the conventional knowledge graph embedding models are inherently transductive, we split the triplets of knowledge graph into train, valid and test sets. Each model is trained on the training set and evaluated on the link  prediction task. We conduct 5 knowledge graph embedding models , including TransE (Bordes et al., 2013), Dist-Mult (Yang et al., 2015), ComplEx (Trouillon et al., 2016), SimplE (Kazemi and Poole, 2018) and Ro-tatE (Sun et al., 2019b). Because their original implementations do not scale to Wikidata5m, we benchmark these methods using the multi-GPU implementation in GraphVite . The performance of link prediction is evaluated in the filtered setting, where test triplets are ranked against all candidate triplets that are not observed in the knowledge graph. We report the standard metrics of Mean Rank (MR), Mean Reciprocal Rank (MRR) and Hits at N (HITS@N). Table 4 shows the benchmarks of popular methods on Wikidata5m.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
In this section, we introduce the experiment settings and experimental results of KEPLER on various NLP and KE tasks.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
In experiments, we choose RoBERTa (Liu et al., 2019c) as our base model and implement our methods in the fairseq framework (Ott et al., 2019) for pre-training. Due to the computing resource limit, we choose RoBERTa BASE architecture and use the released roberta.base 4 parameters to initialize our model.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
In our pre-training procedure, we only use the English Wikipedia corpus to save time and also for a fair comparison with previous knowledgeenhanced PLMs (Zhang et al., 2019; 
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
In this section, we introduce how our KEPLER can be used as a knowledge-enhanced PLM on various NLP tasks and its performance compared with state-of-the-art models.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
Relation classification is an important NLP task that requires models to classify relation types between two given entities from text. We evaluate our model and baselines on two commonlyused datasets: TACRED (Zhang et al., 2017) and FewRel . TACRED covers 42 relation types and contains 106,264 sentences. FewRel is a few-shot relation classification dataset, which has 100 relations and 700 instances for each relation.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
Here we follow the relation extraction finetuning procedure from Zhang et al. (2019), where four special tokens are added before and after entity mentions in the sentence to highlight where the entities are. Then we take the [CLS] output as the sentence representation for classification. Table 5 shows results of various models on TA-CRED, from which we can see that our model Model 5-way 1-shot 5-way 5-shot 10-way 1-shot 10-way 5-shot  Table 6: Accuracies (%) on FewRel dataset. "Proto" indicates Prototypical Networks (Snell et al., 2017) used in . "PAIR" is proposed in Gao et al. (2019)   achieves state-of-the-art on this benchmark. Note that some baselines use the LARGE version of pretrained language models while we still take the BASE architecture. We have gained a large promotion over our base model (RoBERTaBASE) while staying a little bit advanced over other competitive methods (even if they use a LARGE architecture). Our model has also shown strength on FewRel dataset. We use Prototypical Networks (Snell et al., 2017) and PAIR (Gao et al., 2019) as the base frameworks and try out different kinds of pretrained models as encoders. As shown in Table  6, for both frameworks, our models have superior performance over others. We have also compared with current state-of-the-art MTP (Baldini Soares et al., 2019), which outperforms us a little. But note that MTP uses a large version of BERT while we use the base version, and also it carries out a new pre-training task specifically targeting relation extraction, while ours is a general way to combine knowledge and natural language which would benefit all knowledge-related tasks.  
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
Entity typing requires models to classify given entity mentions into pre-defined entity types. For this task, we evaluate all the models on OpenEntity (Choi et al., 2018) following the setting from Zhang et al. (2019), which focuses on nine general entity types. Evaluation results are demonstrated in Table  7. For now we have achieved better results than RoBERTa, and ERNIE and KnowBERT show slightly better results than ours. It is mainly due to that we use different ways of extracting entity representations. KnowBERT adds special tokens before and after the mention and uses the output of the token before the mention as the representation for typing, while ours, for now, directly uses [CLS]. We will try this better way of entity representation in the future.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
In this section, we show how our KEPLER works as a KE model, and evaluate it on our Wikidata5m dataset in inductive setting.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
We do not use the existing KE benchmarks because they are lack of high-quality text descriptions for their entities and they do not have a reasonable data split for the inductive setting.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
We evaluate the generalization ability of our KE-PLER by testing it on the inductive setting in Wiki-data5m (as described in Section 4.2), which re-quires it to produce effective entity embeddings for the unseen entities. The results are shown in Table 8.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
In this paper, we propose KEPLER, a unified model for knowledge embedding and pre-trained language representation. We jointly train the knowledge embedding and language representation objectives on top of the language representation model. Experimental results on extensive tasks demonstrate the effectiveness of our model.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
In the future, we will: (1) Evaluate whether our model can recall factual knowledge with more tasks.
KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation
(2) Try variations of existing models, such as highlighting entity mentions in descriptions or changing knowledge embedding form, to get better understanding of how KEPLER works and bring more promotion for downstream tasks.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Recent language models (LMs) such as BERT  and its successors are remarkable at memorizing knowledge seen frequently during training, however performance degrades over the long tail of rare facts. Given the importance of factual knowledge for tasks such as questionanswering, search, and personal assistants (Bernstein et al., 2012;Poerner et al., 2020;Orr et al., 2020), there has been significant interest in injecting these base LMs with factual knowledge about entities (Zhang et al., 2019;Peters et al., 2019, inter alia.). In this work, we work we propose a simple and effective approach for enhancing LMs with knowledge, called metadata shaping.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Existing methods to capture entity knowledge more reliably, typically use the following steps: first annotating natural language text with entity metadata, and next modifying the base LM model to learn from the tagged data. Entity metadata is obtained by linking substrings of text to entries in a knowledge base such as Wikidata, which stores entity IDs, types, descriptions, and relations. Model modifications include introducing continuous vector representations for entities or auxiliary objectives (Zhang et al., 2019;Peters et al., 2019;Yamada et al., 2020;Xiong et al., 2020;Joshi et al., 2020a;Su et al., 2021). Other methods combine multiple learned modules, which are each specialized to handle fine-grained reasoning patterns or subsets of the data distribution (Chen et al., 2019;Wang et al., 2021).
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
These knowledge-aware LMs have led to impressive gains compared to base LMs on entityrich tasks. That said, the new architectures are often designed by human experts, costly to pretrain and optimize, and require additional training as new entities appear. Further, these LMs may not use the collected entity metadata effectively -Wikidata alone holds over ∼ 100M unique entities, however many of these entities fall under similar categories, e.g., "politician" entities. Intuitively, if unseen entities encountered during inference share metadata with entities observed during training, the LM trained with this information may be able to better reason about the new entities using patterns learned from similar seen entities. However, the knowledge-aware LMs learn from individual entity occurrences rather than learning these shared reasoning patterns. Implicitly learning entity similarities for 100M entities may be challenging since 89% of the Wikidata entities do not appear in Wikipedia, a popular source of unstruc- The FewRel benchmark involves identifying the relation between a subject and object string. The above subject and object are unseen in the FewRel training data and the tuned base LM reflects low attention weights on those words. A base LM trained with shaped data reflects high attention weights on useful metadata words such as "politician". Weights are shown for words which are not stop-words, punctuation, or special-tokens. tured training data for the LMs, at all. 1 We thus ask, to what extent can we match the quality of knowledge-aware LM architectures using the base LM itself? We find that applying some simple modifications to the data at train and test time, a method we call metadata shaping, is surprisingly quite effective. Given unstructured text, there are several readily available tools for generating entity metadata at scale (e.g., Manning et al. (2014);Honnibal et al. (2020)), and knowledge bases contain entity metadata including type tags (e.g., Barack Obama is a "politician") and descriptions (e.g., Barack Obama "enjoys playing basketball"). Our method entails explicitly inserting retrieved entity metadata in examples as in Figure  1 and inputting the resulting shaped examples to the LM. Our contributions are:
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Simple and Effective Method We propose metadata shaping and demonstrate its effectiveness on standard benchmarks that are used to evaluate knowledge-aware LMs. Metadata shaping, with simply an off-the-shelf base LM, exceeds the base LM trained on unshaped data by by an average of 4.3 F1 points and is competitive to state-of-theart methods, which do modify the LM. Metadata shaping thus enables re-using well-studied and optimized base LMs (e.g., ).
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Tail Generalization We show that metadata shaping improves tail performance -the observed gain from shaping is on average 4.4x larger for the slice of examples containing tail entities than for the slice containing popular entities. Metadata establish "subpopulations", groups of entities sharing similar properties, in the entity distribution (Zhu et al., 2014;Cui et al., 2019;Feldman, 2020). For example on the FewRel benchmark (Han et al., 2018), "Daniel Dugléry" (a French politician) appears 0 times, but "politician" entities in general appear > 700 times in the task training data. Intuitively, performance on a rare entity should improve if the LM has the explicit information that it is similar to other entities observed during training.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Explainability Existing knowledge-aware LMs use metadata (Peters et al., 2019;Alt et al., 2020), but do not explain when and why different metadata help. Inspired by classic feature selection techniques (Guyon and Elisseeff, 2003), we conceptually explain the effect of different metadata on generalization error.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We hope this work motivates further research on addressing the tail challenge through the data. 2
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
This section introduces metadata shaping, including the set up and conceptual framework.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
The goal of metadata shaping is to improve tail performance using properties shared by popular and rare examples (e.g., the unseen entity "Daniel Dugléry" and popular entity "Barack Obama" are both "politicians"). This work explores how to effectively provide these properties to popular transformer models. Tail entities are those seen < 10 times during training and head entities are seen ≥ 10 times, consistent with Orr et al. (2020); Goel et al. (2021).
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Metadata are easily and scalably sourceable using off-the-shelf models such as those for named entity (NER, NEL) or part-of-speech (POS) tagging (Manning et al., 2014;Honnibal et al., 2020), heuristic rules, and knowledge bases (KBs) (e.g., Wikidata, Wordnet (Miller, 1995), domain-specific KBs (Bodenreider, 2004), and product KBs (Krishnan, 2018)). KBs often provide high tail coverage -e.g., a product KB will contain metadata for both popular and unpopular products.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Many prior works annotate text with metadata and in our setting, instead of using predefined feature schemas (Marcus et al., 1993;Mintz et al., 2009, inter alia.), we consider using an unrestricted set of metadata, including entity unstructured descriptions. Importantly, knowledge-aware LMs have attracted significant recent interest and dataoriented approaches have not been demonstrated as a compelling alternative, the aim of this work.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Let input x ∈ X and label y ∈ Y, and consider the classification dataset D D D = {(x i , y i )} n i=1 of size n. Let m ∈ M denote a metadata tag and let M (x i ) M (x i ) M (x i ) be the set of metadata collected for example x i . A shaping function f s : X → X s accepts an original example x i ∈ X and produces a shaped example s i ∈ X s by inserting a subset of M ( This work uses the following representative metadata shaping functions for all tasks to insert a range of coarse-grained signals associated with groups of examples to fine-grained specific signals associated with individual examples:
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Categorical tokens establish subpopulations of entities (e.g., Dugléry falls in the coarse grained category of "person" entities, or finer grained category of "politician" entities). NER and POS tags are coarse grained categories, and knowledge bases contain finer-grained categories (i.e., entity types and relations). Categories are consistent and frequent compared to words in the original examples.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Description tokens give cues for rare entities and alternate expressions of popular entities (e.g., Dugléry is a "UMP party member"). Descriptions are likely unique across entities, and can be viewed as the finest-grained category for an entity.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Next we want to understand if inserting m ∈
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We measure the generalization error of the classification modelp φ using the cross-entropy loss:
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Let Pr(y|x i ) be the true probability of class y ∈ Y given x i . Example x i is composed of a set of patterns K i K i K i (i.e., subsets of tokens in x i ). We make the assumption that a pattern k ∈ K i K i K i is a useful signal if it informs Pr(y|x i ). We thus parametrize the true distribution Pr(y|x i ) using the principle of maximum entropy (Berger et al., 1996):
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
where λ k represents learned parameters weighing the contributions of patterns (or events) k and Z(x i ) is a partition function that ensures Pr(y|x i ) represents a probability distribution. Therefore when evaluatingp φ , achieving zero cross-entropy loss between the true probability Pr(y|k) and the estimated probabilityp φ (y|k), for all k, implies zero generalization error overall.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Unseen Patterns Our insight is that for a pattern k that is unseen during training, which is common in entity-rich tasks, 3 the class and pattern are in-
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
With the assumption of a well-calibrated model and not considering priors from the base LM pretraining stage, 4 this probability isp φ (y) = 1 |Y| for y ∈ Y. Plugging inp φ (y) = 1 |Y| , the cross-entropy loss between Pr(y|k) andp φ (y|k) is Pr(k) log |Y |. Our idea is to effectively replace k with another (or multiple) shaped pattern k , which has nonuniformp φ (y|k ) and a lower cross-entropy loss with respect to Pr(y|k ), as discussed next. Compute pmi(y, m) for y ∈ Y. 5: end for 6: for each class y ∈ Y over D train do 7:
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Compute frequency f y . 8: end for 9:
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
10: Select Metadata for Sentence 11: Input:
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Compute r y = 2 pmi(m,y) f y for y ∈ Y.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Normalize r y values to sum to 1.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Compute entropy H m over r y for y ∈ Y. 17: end for
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
For a rare pattern (e.g., a mention of a rare entity in x i ) k, if an associated pattern k m (e.g., a metadata token for the rare entity) occurs non-uniformly across classes during training, then the cross-entropy loss betweenp φ (y|k m ) and Pr(y|k m ) is lower than the cross-entropy loss betweenp φ (y|k) and Pr(y|k). If k m shiftsp φ (y|x i ) usefully, performance ofp φ should improve.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We can measure the non-uniformity of k m across classes using the conditional entropyĤ(Y|k).
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
For non-uniformp φ (y|k m ), the conditional entropy decreases. Broadly, we connect the benefit of using different metadata, which are inputs both to existing knowledge aware LMs and our approach, to classical methods (Guyon and Elisseeff, 2003) -we seek the metadata providing the largest information gain. Next we discuss the practical considerations for selecting metadata.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
M (x i ) -categories can range from coarse-grained (e.g., "person") to fine-grained (e.g., "politician" or "US president") and there are intuitively many ways to describe entities. Since certain metadata may not be helpful for a task, and popular base LMs do not scale very well to long sequences (Tay et al., 2020;Pascanu et al., 2013), it is important to understand which metadata to use for shaping.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We want to select k m with non-uniform p φ (y|k m ) across y ∈ Y, i.e. with lowerĤ(Y|k m ). Conditional probability Pr(y|k m ) is defined as:
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
where we recall that the pointwise mutual information pmi(y, k m ) is defined as log Pr(y,km) Pr(y) Pr(km) . The pmi compares the probability of observing y and k m together (the joint probability) with the probabilities of observing y and k m independently. Class-discriminative metadata reduceĤ(Y|k).
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Directly computing the resulting conditional probabilities after incorporating metadata in D D D is challenging since the computation requires considering all patterns contained in all examples, generated by including m. Instead we use simplistic proxies to estimate the information gain. In Algorithm 1, we focus on the subset of K s i K s i K s i containing individual metadata tags m, and compute the entropy overp φ (y|m) for y ∈ Y. Simple extensions to Algorithm 1, at the cost of additional computation, would consider a broader set of k m (e.g., n-grams containing m for n > 1), or iteratively select tokens by considering the correlations in the information gain between different metadata tags.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
In this section, we demonstrate that metadata shaping is general and effective.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We evaluate on standard entity-typing and relation extraction benchmarks used by baseline methods. Entity typing involves predicting the the applicable types for a given substring in the input example from a set of output types. We use OpenEntity (9 output types) (Choi et al., 2018) for evaluation. Relation extraction involves predicting the relation between the two substrings in the input example, one representing a subject and the other an  Table 1: Test scores on standard relation extraction and entity-typing tasks. "Ours (Base LM)" is metadata shaping. All methods use the same base LM (BERT-base) and external information (Wikipedia) for consistent comparison. A dash ("-") indicates the baseline method did not report scores for the task. object. We use FewRel (80 output relations) and TACRED Revisited (42 output relations) for evaluation (Han et al., 2018;Zhang et al., 2017;Alt et al., 2020). While metadata shaping is generally applicable to classification tasks, our objective in this work is to compare architectural versus dataoriented methods of injecting knowledge, so we focus on benchmarks that are popular in the literature on knowledge-aware LMs.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Model We fine-tune a BERT-base model on metadata shaped data for each task, taking the pooled [CLS] representation and using a linear prediction layer for classification . We use cross-entropy loss for FewRel and TACRED and binary-cross-entropy loss for OpenEntity. All test scores are reported at the epoch with the best validation score and we use the scoring implementations released by (Zhang et al., 2019). Additional training details are provided in appendix A.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We collect entity metadata from Wikidata for our evaluations, a compelling choice as several works successfully improve tail performance in industrial workloads using the knowledge base (e.g., Orr et al. (2020)) We use the state-of-the-art pretrained entity-linking model from Orr et al. (2020) to link the text in each task to an October 2020 dump of Wikidata. We use Wikidata and the first sentence of an entity's Wikipedia page to obtain descriptions. Additional details are in Appendix A. For certain examples in the tasks, there are no linked entities in the text (e.g., several subject or object entities are simply pronouns or dates). Table 3 gives statistics for the number of examples with available of metadata for each task. Metadata tags are selected by Algorithm 1.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
While the metadata annotation methods have their own failure rates, our baselines also use entity linking as the first step (Zhang et al., 2019, inter alia.) with the same exposure to failures. All the same, we seek methods that are flexible to errors that arise in natural data.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Prior work proposes various knowledge-aware LMs, which are currently the state-of-the-art for the evaluated tasks. ERNIE, (Zhang et al., 2019) LUKE (Yamada et al., 2020), KEPLER , CokeBERT (Su et al., 2021), and WKLM (Xiong et al., 2020) introduce auxilliary loss terms and require additional pretraining. Prior approaches also modify the architecture for example using alternate attention mechanisms (KnowBERT (Peters et al., 2019), K-BERT , LUKE) or training additional transformer stacks to specialize in knowledge-based reasoning (K-Adapter (Wang et al., 2021)). E-BERT (Poerner et al., 2020) does not require additional pretraining and uses entity embeddings which are aligned to the word embedding space. In Table 1, we compare to methods which use the same base LM, BERT-base, and external information resource, Wikipedia, for consistency.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We simply use an off-the-shelf BERT-base LM , with no additional pretraining and fine-tuned on shaped data to exceed the BERT-base LM trained on unshaped data by 5.3 (FewRel), 4.7 (TACRED), and 3.0 (OpenEntity) F1 points. Metadata shaping is also competitive with SoTA baselines which do modify the BERTbase LM. Results are shown in Table 1. Table 3 reports the availability of metadata for each task.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We observe that metadata shaping is effective both when most task examples have available metadata (e.g., FewRel) and when metadata tags are sparse (e.g., on OpenEntity only 30% of examples have available metadata), analyzed further in Section 4. We further note that the performance of our method is not sensitive to grammatical choices around how the metadata tags are inserted through ablations provided in Appendix B.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
For the baselines, we give reported numbers when available, Su et al. (2021) reports two of the KnowBERT-Wiki and all K-BERT results, and we obtain remaining numbers using the code released by baseline work as detailed in Appendix A.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Here we study the following key questions for effectively using metadata shaping: Section 4.1 What are the roles of different varieties of metadata? Section 4.2 What are the effects of metadata shaping on slices concerning tail versus popular entities?
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Metadata Effects Class-discriminative metadata correlates with reduced model uncertainty. High quality metadata, as found in Wikidata, results in improved classification performance.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
To investigate the effects of metadata on model uncertainty, we compute the entropy ofp φ softmax scores over the output classes as a measure of uncertainty, and compute the average across test set examples. Lower uncertainty is correlated with improved classification F1 (See Figure 2 (Left)).
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We compute pmi scores for inserted metadata tokens as a measure of class-discriminativeness. We rank individual tokens k by pmi(y, k) (for task classes y), computed over the training dataset. On FewRel, for test examples containing a top-20 pmi word for the gold class, the accuracy is 27.6% higher when compared to the slice with no top-20 pmi words for the class. Notably, 74.1% more examples contain a top-20 pmi word for their class when pmi is computed on shaped data vs. unshaped training data.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Metadata Selection Simple information theoretic heuristics are effective for selecting metadata, despite the complexity of the underlying contextual embeddings.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We apply Algorithm 1, which ranks metadata tags by their provided information gain, to select metadata tags for the tasks. Given x i with a set  M (x i ) M (x i ) M (x i ) of metadata tags, our goal is to select n to use for shaping. We compare four selection approaches: using the highest ("High Rank") and lowest ("Low Rank") ranked tokens by Algorithm 1, random metadata from M (x i ) M (x i ) M (x i ) ("Random"), and the most popular metadata tokens across the union of M (x i ) M (x i ) M (x i ), ∀x i ∈ D train D train D train ("Popular"), selecting the same number of metadata tags per example for each baseline. We observe that High Rank consistently gives the best performance, evaluated over three seeds, and note that even Random yields decent performance vs. the BERT-baseline, indicating the simplicity of the method (Table 2).
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Considering the distribution of selected category tokens under each scheme, the KL-divergence between the categories selected by Low Rank vs. Popular is 0.2 (FewRel), 4.6 (OpenEntity), while the KL-divergence between High Rank vs. Popular is 2.8 (FewRel), 2.4 (OpenEntity). Popular tokens are not simply the best candidates; instead, Algorithm 1 selects discriminative metadata.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
For OpenEntity, metadata are relatively sparse, so categories appear less frequently in general and it is reasonable that coarse-grained types have more overlap with High Rank. For e.g., "business" is in the top-10 most frequent types under High Rank, Figure 2: Test F1 forp φ (no additional pretraining) vs. average entropy ofp φ softmax scores (Top) and vs. perplexity of a modelp θ (w/ pretraining) (Bottom).p φ and p θ use the same shaped training data. Each point is a different metadata shaping scheme (median over 3 Random Seeds): for R0 all inserted tokens are true tokens associated with the entity in the KB. For RX, X true metadata tokens are replaced by random (noise) tokens from the full vocabulary. For each point, the total number of metadata tokens is constant per example.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
while "non-profit" (occurs in 2 train examples) is in the top-10 most frequent types for Low Rank. Metadata tokens overall occur more frequently in FewRel (See Table 3), so fine-grained types are also quite discriminative. The most frequent category under Low Rank is "occupation" (occurs in 2.4k train examples), but the top-10 categories under High Rank are finer-grained, e.g. "director" and "politician" (each occurs in > 300 train examples).
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Task Agnostic Metadata Effects Using metadata correlates with reduced task-specific LM uncertainty. We observe shaping also correlates with reduced LM uncertainty in a task-agnostic way.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We perform additional masked language modeling (MLM) over the shaped task training data using an off-the-shelf BERT-MLM model to learn model p θ . We minimize the following loss function and evaluate the model perplexity on the task test data:
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
where I is the masked token distribution and s m i is the masked token at position i in the shaped sequence s m . 5 Through minimizing the MLM loss, p θ learns direct dependencies between tokens in the data (Zhang and Hashimoto, 2021). In Figure 2 (Right), we observe a correlation between reduced perplexity forp θ , and higher downstream performance forp φ across multiple tasks, both using the same training data. Overall, shaping increases the likelihood of the data, and we observe a correlation 5 We use the Hugging Face implementation for masking and fine-tuning the BERT-base MLM . between the intrinsic perplexity metric and the extrinsic downstream metrics as a result of the same shaping scheme. Metadata Noise We hypothesize that noisier metadata can provide implicit regularization. Noise arises from varied word choice, word order, and blank noising.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Feature noising (Wang et al., 2013) is effective to prevent overfitting and while regularization is typically applied directly to model parameters, Xie et al. (2017); Dao et al. (2019) regularize through the data. We hypothesize that using metadata with diverse word choice and order (e.g., entity descriptions) and blank noising (e.g., by masking metadata tokens), can help reduce overfitting, and we provide initial empirical results in Appendix B.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Section 3 shows the overall gain from shaping. We now consider fine-grained slices of examples containing head vs. tail entities and observe gains are 4.4x larger on the tail slice on average (Figure 3). 6
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Subpopulations Metadata are helpful on the tail as they establish subpopulations.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We hypothesize that if a pattern is learned for an entity-subpopulation occurring in the training data, the model may perform better on rare entities that also participate in the subpopulation, but were not individually observed during training. On FewRel, we take the top-20 TF-IDF words associated with each category signal during training as linguistic cues captured by the model for the category subpopulation, consistent with Goel et al. (2021). For example, "government" is in the top-20 TF-IDF words for the "politician" entity category. At test time, we select the slice of examples containing any of these words for any of the categories inserted in the example. The performance is 9.0/3.5 F1 points higher on examples with unseen subject/object entities with vs. without a top-20 TF-IDF word for a subject/object category.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
For popular entities the LM can learn entity-specific patterns well, and be mislead by subpopulation-level patterns corresponding to metadata.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Although we observe overall improvements, here we examine the effect of metadata on the popular entity slice within our conceptual framework.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Let p be a popular pattern (i.e., entity mention) in the training data, and let m be a metadata token associated with p. Intuitively, the LM can learn entity-specific patterns from occurrences of p, but coarse-grained subpopulation-level patterns corresponding to m. If m and p are class-discriminative for different sets of classes, then m can mislead the LM. To evaluate this, consider subject and object entity spans p ∈ P seen ≥ 1 time during training. For test examples let Y p be the set of classes y for which there is a p ∈ P in the example with pmi(y, p) > 0, and define Y m as the classes y for which there is a metadata token m with pmi(y, m) > 0 in the example. The examples where Y p = ∅, Y m = ∅, and Y p contains the true class, but Y m does not, represents the slice where metadata can mislead the model. On this slice of FewRel, the gain from the shaped model is 2.3 F1 points less than the gain on the slice of all examples with Y p = ∅ and Y m = ∅, supporting our intuition.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
An example entity-specific vs. subpopulationlevel tension in FewRel is: p = "Thames River" is class-discriminative for y ="located in or next to body of water", but its m ="river" is classdiscriminative for y ="mouth of the watercourse".
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We propose metadata shaping to improve tail performance. The method is a simple and general baseline that is competitive with SoTA approaches for entity-rich tasks. We empirically show that the method improves tail performance and explain why metadata can reduce generalization error. While this work focused on entity-rich tasks, metadata shaping is not limited to this setting. Broadly, we hope this work motivates further research on understanding how to effectively program LMs with useful and readily available side information. While modifying the LM architecture to encode the information has been a popular approach, modifying the data is a simple and effective alternative.  
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We report the test score at the epoch with the highest validation score. For the results in Table 1, we evaluated the number of metadata tokens to insert, whether place the tokens directly following or at the end of the example, and whether to use blank noising on the metadata tokens. Metadata tokens are ranked by Algorithm 1.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We use up to 20 metadata categories per subject and object on FewRel, up to 25 metadata categories per subject on OpenEntity, and up to 5 metadata categories per subject and object on TACRED. Note that categories (e.g., "United States federal executive department") can include multiple tokens, selecting these maximum values by grid search. For FewRel and OpenEntity, we insert metadata tokens directly after the corresponding entity mention, and for TACRED, we inserted all metadata at the end of the example. For OpenEntity we randomly mask 10% of metadata tokens at training time as implicit regularization, and for relation extraction, we use no blank noising. The impact of position and blank noising are further discussed in Appendix B.3.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We produce numbers for key baselines which do not report for the benchmarks we consider, using provided code. 8 9
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
• We produce numbers for KnowBERT-Wiki on TACRED-Revisited using a learning rate of 3e − 5, β 2 = 0.98, and choosing the best score for epochs ∈ 1, 2, 3, 4 and the remaining provided configurations.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
• We produce numbers for ERNIE on TACRED-Revisited using the provided training script and configurations they use for the original TACRED task.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
In Table 4 we report the same experiment conducted in Section 4.1, for all benchmark tasks considered in this work. Each point represents the median test score over 3 random seeds.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Noisier metadata appear to provide implicit regularization. Noise arises from varied word choice and order, as found in entity descriptions, or blank noising (i.e. random token deletion).
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Here we provide initial empirical results. Blank noising (Xie et al., 2017) by randomly masking 10% of inserted metadata tokens during training leads to a consistent boost on OpenEntity: 0.1 ("High Rank"), 0.5 ("Popular"), 0.5 ("Low Rank") F1 points higher than the respective scores from Table 2 over the same 3 random seeds. We observe no consistent benefit from masking on FewRel. Since metadata are sparsely available for OpenEntity examples, we hypothesize that blank noising of the category tokens can prevent overreliance on the signal. Future work could investigate advanced masking strategies, for example masking discriminative words in the training data.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Descriptions use varied word choice and order vs. category metadata. 10 To study whether shaping with description versus category tokens lead the model to rely more on metadata tokens, we consider two shaping schemes that use 10 metadata tokens: 10 category tokens and 5 category, 5 description, where the categories are randomly selected. We observe both give the ∼same score Benchmark R 2 FewRel 0.985 TACRED 0.782* OpenEntity 0.956 Table 4: Correlation (R 2 ) between test F1 ofp φ (no additional pretraining) vs. perplexity of the independent modelp θ (w/ additional pretraining) for three tasks, using the procedure described in Figure 2. *Without one outlier corresponding to shaping with all random tokens (R 2 = 0.02 with this point).
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
on FewRel, 89.8 F1 and 89.5 F1, and use models trained with these two schemes to evaluate on test data where 10% of metadata tokens per example are randomly removed. Performance drops by 1.4 F1 for the former and 1.0 F1 for the latter.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We also analyze the degree of sensitivity of metadata shaping to how the metadata are inserted in examples (e.g., special tokens, the number of metadata tokens, and position).
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Boundary Tokens Designating the boundary between original tokens in the example and inserted metadata tokens improves model performance.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Inserting boundary tokens (e.g., "#") in the example, at the start and end of a span of inserted metadata, consistently provides a boost across the tasks. Comparing performance with metadata and boundary tokens to performance with metadata and no boundary tokens, we observe a 0.7 F1 (FewRel), 1.4 F1 (OpenEntity) boost in our main results. We use boundary tokens for all results in this work.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Task Structure Tokens designate relevant entities in the examples (e.g., "[START_SUBJECT]" and "[END_SUBJECT]"). With no other shaping, inserting these tokens provides a 26.3 (FewRel), 24.7 (OpenEntity) F1 point boost vs. training the BERT model without task structure tokens. These tokens are already commonly used.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
Token Insertion We observe low sensitivity to increasing the context length and to token placement (i.e., inserting metadata directly-following the entity-span vs at the end of the sentence).
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We evaluate performance a the maximum number of inserted tokens per entity, n, increases. 11 11 Per subject and object entity for FewRel, and per main entity for OpenEntity. I.e., n = 10 for FewRel yields a maximum of 20 total inserted tokens for the example.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
We insert metadata tokens in a random order (to control for the effect of different metadata having different levels of class-discriminativeness) and observe that for FewRel, n ∈ {1, 5, 10, 20} gives {85.4,86.4,87.6,88.5} test F1. On OpenEntity,n ∈ {1,5,10,20,40} gives {74.9,75.7,74.8,74.5,75.8} test F1. Overall performance changes gracefully with n and we observe low sensitivity to longer contexts.
Metadata Shaping: A Simple Approach for Knowledge-Enhanced Language Models
The benefit of inserting metadata directlyfollowing the entity span vs at the end of the example differed across tasks (e.g., for TACRED, placement at the end performs better, for the other tasks, placement directly-following performs better), though the observed difference was small. In Section 4, tokens are inserted directly-following the relevant entity span for all tasks.
Explaining Toxic Text via Knowledge Enhanced Text Generation
The toxic speech detection and classification problem has seen increasing interest in recent years. However, it is not only important for AI agents to recognize and classify toxic speech, but to also explain why it is toxic. For instance, debiasing methods that use information about toxic language may benefit from additional information given by detailed explanations of toxicity in text (Ma et al., 2020). Furthermore, detailed explanations of toxicity may facilitate human interaction with toxicity detection systems (Rosenfeld and Richardson, 2019). They can also help humans who work with toxicity classifiers use more information about the input when making decisions about toxic speech. To elucidate, consider the following offensive joke: "What type of punch do you use against a kindergartener? A sandy-hook.". While the literal text is not toxic, the implied meaning is offensive, particularly to those affected by school shootings. An AI agent capable of generating the implied meaning could thus provide additional information to downstream actors. Note that, we use the term biased and toxic interchangeably in this work.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Existing work largely addresses the problem of detecting and classifying toxic speech (Waseem and Hovy, 2016;Founta et al., 2018;Davidson et al., 2017). As mentioned earlier, explanations of toxicity can help with downstream tasks such as debiasing or decision making by humans, thus there has been increasing demand for explainable machine learning classifiers (Ribeiro et al., 2016;Došilović et al., 2018). Recent work around explainable toxicity classification introduced Social Bias Frames (Sap et al., 2020), a formal framework which combines explanations of toxicity along with toxicity classifications along multiple dimensions. However the explanations generated from the current state-of-the-art methods tend to be generic, without much detail. For instance, explanations may focus on certain toxic components of the input but ignore others, or include irrelevant stereotypes about the minority group affected.
Explaining Toxic Text via Knowledge Enhanced Text Generation
To fill this gap, our work proposes to leverage different types of knowledge to provide rich context and background for toxicity explanation. Specifically, we introduce a novel framework to utilize three distinct knowledge sources. Prior work (Yu et al., 2022) divides knowledge broadly into internal and external knowledge, where internal knowledge is knowledge embedded in the input text, and external knowledge is derived from sources outside the input. Building upon these, we leverage expert knowledge that comes from high-quality expert annotations of the input, and explicit knowledge from knowledge graphs and bases, as such symbolic knowledge can provide relevant information to the output text (Yu et al., 2022;Mou et al., 2016). While knowledge graphs and bases deterministically retrieve and restructure knowledge from raw text sources, large pretrained generative models are 811 found to be effective in outputting useful knowledge in a probabilistic manner, complementing the expert and explicit knowledge (Razniewski et al., 2021). To this end, we also include implicit knowledge models which source knowledge from large pretrained text generation models. We further build a family of mixture models, MIXGEN, to synthesize knowledge from all three sources, as shown in Figure 1. To sum up, our contributions are twofold: (1) We leverage three different sources of knowledge, and further combine them using simple yet effective mixture models to explain toxic text.
Explaining Toxic Text via Knowledge Enhanced Text Generation
(2) We show that our models outperform prior stateof-the-art baselines and generate more detailed explanations.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Internal knowledge includes knowledge that is available within the input. For instance, Latent Dirichlet Allocation (LDA) (Blei et al., 2003) can learn topics from inputs, which can then be incorporated into text generation models (Cao et al., 2015;Guo et al., 2020). Keywords may also be extracted from input text using techniques like TF-IDF, PMI or independent classifiers. In one such work, Song et al. (2019) extract emotion oriented keywords using an independent emotion classifier to enhance dialogue generation. Similarly, Mou et al. (2016) use PMI to find relevant keywords for short text conversation. Forbes et al. (2020) develop conceptual formalisms that rely on annotations about the input to generate text. In this work, we denote the use of independent annotations on input to enhance text generation as expert knowledge, as such annotations often come from human experts.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Knowledge graphs and bases are commonly used as a form of external knowledge. Zhang et al. (2019) use knowledge graph embeddings to model conversation flow, while Guan et al. (2020) use triples extracted from knowledge graphs to enhance story generation. Finally, Lian et al. (2019) develop probabilistic mechanisms to select knowledge from knowledge bases for response generation. Knowledge from conventional sources are determinstically created, in that they simply restructure raw text and store them in a knowledge base or graph. We refer to this type of external knowledge as explicit knowledge. On the other hand, there has been increasing interest in the use of large pretrained generative models as a source of knowledge. Heinzerling and Inui (2021) argue that large pretrained models can in fact serve as knowledge bases, while Davison et al. (2019) argue that pretrained models can accurately assess the validity of knowledge mined from raw text. While pretrained models are also trained on raw texts, similar to knowledge bases and graphs, they draw from this knowledge probabilistically and thus are a distinct approach. We denote this type of external knowledge as implicit knowledge.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Prior work around toxicity understanding mainly focuses on detection (Schmidt and Wiegand, 2017). Early approaches include using n-grams (Waseem and Hovy, 2016;Sood et al., 2012;Perera and Fernando, 2021) as well as word clustering (Xiang et al., 2012;Zhong et al., 2016). Recently, knowledge enhanced approaches have also been used for toxicity detection. For instance, Dinakar et al. (2012) use ConceptNet to detect anti-LGBT bullying. The use of meta-information, such as information about the user (Dadvar et al., 2012), has proven to be useful, depending on the type of information used. Sap et al. (2020) use Social Bias Frames to produce both toxicity classifications and explanations of toxicity. Similarly, our approach attempts to explain toxicity by leveraging different sources of knowledge to provide more context and grounding for the models to generate explanations. Different from many prior works, we synthesize these diverse knowledge sources in a unified framework to utilize the unique contribution from each 812 individual knowledge source.
Explaining Toxic Text via Knowledge Enhanced Text Generation
This section presents our selected three different types of knowledge-expert knowledge, explicit knowledge and implicit knowledge, and our MIX-GEN models for toxicity explanation.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Expert knowledge is sourced from annotations of the input. For instance, in the Social Bias Frames dataset (Sap et al., 2020), such expert knowledge include human judgements towards the lewdness, offensiveness, intent to offend, and group targeted categories. This type of expert knowledge provides useful insights and heuristics for the toxicity explanation task, if they are available.
Explaining Toxic Text via Knowledge Enhanced Text Generation
We incorporate expert knowledge into the generation process using the join embedding technique (Pryzant et al., 2020) along with toxicity classification models. The join embedding architecture uses attention weights from the toxicity classifiers to inform the text generation model about parts of the input post relevant to toxicity classification, thus providing a heuristic for the related toxicity explanation task. Formal details of the architecture can be found in Appendix A.1.
Explaining Toxic Text via Knowledge Enhanced Text Generation
We denote these models with the naming convention, "EXPERT [FEATURE]", where "[FEATURE]" is the categorical variable we use for the join embedding. We replace "[FEATURE]" with "ALL" when we train on all features.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Explicit knowledge is sourced from some knowledge base or graph. Common sources include Con-ceptNET, DBpedia, WikiData, etc. (Auer et al., 2007;Vrandečić and Krötzsch, 2014). We opt to use ConceptNet (Speer et al., 2017), since it contains commonsense knowledge (Liu and Singh, 2004). Commonsense knowledge incorporates everyday concepts, especially knowledge regarding social groups and situations.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Following Chang et al. (2020), given a BART model and an input, we extract ranked triples related to the input post and keep the top k triples per post where we vary the k ∈ {3, 5, 10, 15, 20, 25}. We experiment with both concatenation and attention based methods to incorporate the top k triples, but settle on a concatenation based approach due to its simplicity and the lack of performance gains from the attention based approach. Results and analysis for both the concatenation and attention based approaches are provided in Appendix 12. We denote these models with the naming convention, "EXPLICIT (K)", where K denotes the number of triples used.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Implicit knowledge is obtained from some text generator, such as a large pretrained generative model. Prior work such as Heinzerling and Inui (2021), argue that large pretrained models can in fact serve as knowledge bases. Implicit knowledge grants models a probabilistic view of external raw text sources related to a given scenario or input, since generative models tend to generate based on statistical correlations found in their training corpora (Razniewski et al., 2021).
Explaining Toxic Text via Knowledge Enhanced Text Generation
To use implicit knowledge, we first train a BART model to generate the target minority group from the input post. Following Sheng et al. (2019), we use the predicted target minority corresponding to each input post and a set of prompts to induce biased prompt completions from GPT models. We may use multiple prompts per input post, where the number of prompt completions generated is governed by a hyperparameter, k. Then we train an independent BART model to generate these biased prompt completions, given the origin input post. This BART model is then retrained to produce the implied stereotype, given the input post. Again, we provide a formal description in Appendix A.3.
Explaining Toxic Text via Knowledge Enhanced Text Generation
We denote these models with "IMPLICIT [GPT|GPT-2] (K)", where "GPT" and "GPT-2" correspond to the model used for prompt completion, while K corresponds to the number of biased prompt completions generated per input post.
Explaining Toxic Text via Knowledge Enhanced Text Generation
We introduce a simple and effective approach to combine all three knowledge sources as input for our MIXGEN family of models, and generate the final stereotype by integrating complementary knowledge from these sources. In our design, we take inspiration from Mixture of Experts models (Masoudnia and Ebrahimpour, 2012), which combine base expert model outputs into a final output using a gating mechanism. Here, we rely on attention mechanisms over the knowledge informed model outputs to serve as the gating mechanism.
Explaining Toxic Text via Knowledge Enhanced Text Generation
We build two variants. The first variant is called MIXGEN CONCAT which uses concatenation to  combine outputs from the knowledge informed models, as shown in Figure 2. The second variant is called MIXGEN MULTIVIEW which uses views to perform self attention over outputs of the knowledge informed models (Chen and Yang, 2020). Since the BART model already uses self attention over input tokens, we experiment with the MultiView architecture to see whether the additional self attention mechanisms of the MultiView model causes changes in performance.
Explaining Toxic Text via Knowledge Enhanced Text Generation
For MIXGEN CONCAT, suppose we have k trained models, M 1 , . . . , M k , each trained to produce the implied stereotype given the input post, and each informed by one of the aforementioned knowledge types.
Explaining Toxic Text via Knowledge Enhanced Text Generation
We con- 
Explaining Toxic Text via Knowledge Enhanced Text Generation
and v 3i are view tokens. Here, v 1i is always the first token in the view string and v 2i and v 3i surround M i 's output. We configure k such views (one for each model) and pass each into the BART MultiView model as a set of views corresponding to the original input post. The BART MultiView model is then trained to produce the corresponding output stereotype. For details on the MultiView architecture, please see Chen and Yang (2020).
Explaining Toxic Text via Knowledge Enhanced Text Generation
We utilize the BART encoder decoder framework throughout (Lewis et al., 2020). We use batch gradient descent when training. For a batch B with padded input sequences X i of length N s and corresponding padded target sequences Y i of length N t , along with knowledge K i from some knowledge  source, we minimize cross entropy loss:
Explaining Toxic Text via Knowledge Enhanced Text Generation
4 Experimental Setup
Explaining Toxic Text via Knowledge Enhanced Text Generation
We conduct our experiments on the SBIC dataset (Sap et al., 2020) and the Implicit Hate dataset (ElSherief et al., 2021). The SBIC dataset contains an input post, toxicity annotations and free text annotations of the implied stereotype. We work with the input post and the the implied stereotype. The Implicit Hate dataset (ElSherief et al., 2021) contains free text annotations of the implied stereotype. Dataset statistics are provided in Table 2.
Explaining Toxic Text via Knowledge Enhanced Text Generation
We compare our models with BART, and state of the art baselines from Sap et al. (2020):
Explaining Toxic Text via Knowledge Enhanced Text Generation
• GPT: Following Sap et al. (2020), we train the GPT pretrained model from huggingface to generate the toxicity classifications, the Target Minority, and the Implied Stereotype as a string, when prompted with the input post.
Explaining Toxic Text via Knowledge Enhanced Text Generation
• GPT-2: We train with the same setting as the GPT Baseline, but use the GPT-2 pretrained model from huggingface.
Explaining Toxic Text via Knowledge Enhanced Text Generation
• BART: We train a standard pretrained BART model to generate the implied stereotype when given the input post.
Explaining Toxic Text via Knowledge Enhanced Text Generation
We use BLEU (Papineni et al., 2002), ROUGE-L (Lin, 2004) and BERTScore (Zhang* et al., 2020) to evaluate our models and take the maximum score for each hypothesis over all of the corresponding references. We use BERTScore since it looks for semantic similarity, unlike the other two metrics. While we acknowledge it as a limitation, we ultimately do not use human evaluation for multiple reasons. First, the generated stereotypes are minimal in length compared to other text generation tasks. Moreover, we perform manual analyses of the results in Section 5. Second, we are following precedent set by prior work (Sap et al., 2020). Finally, we want to minimize annotator exposure to harmful content.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Our results on both dev and test are described in Table 3. Here we focus on dev since both sets of results track similar trends. We observe that the MIXGEN models outperform all other models. After MIXGEN, the model using Implicit Knowledge sources perform best. These are followed by the model using Explicit Knowledge, in turn followed by model using Expert Knowledge. Both models with explicit and with implicit knowledge outperform expert language models. The models using implicit knowledge tend to perform best overall (BLEU: from 0.650 to 0.683, ROUGE-L: from 0.624 to 0.659, BERTScore: from 0.759 to 0.800). This is likely because implicit knowledge is less structured and hence easier to induce bias from (Petroni et al., 2019). On the other hand, while our source (ConceptNET) for explicit knowledge may be biased (Mehrabi et al., 2021), retrieved stereotypes are often mixed with general, unbiased facts.
Explaining Toxic Text via Knowledge Enhanced Text Generation
The MIXGEN models outperform every other model. This makes intuitive sense since MIXGEN synthesizes multiple types of knowledge. Unexpectedly, MIXGEN MULTIVIEW model does not improve performance (the absolute difference is within 0.002 across all scores) over MIXGEN CONCAT. This is likely due to the fact that the MultiView model was intended to capture metasequences in text (Chen and Yang, 2020), whereas in our setting the input is not sequential. We also note that the MIXGEN models perform better than the source models, despite their input being sourced from the source models. Thus, despite differing performance, models from different knowledge sources are likely providing some distinct and complementary information.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Results on the implicit hate corpus (ElSherief et al., 2021) are given in  Table 3: We report performance of baseline models (first three rows) and our representative models from each knowledge source. A superscript of * indicates statistically significant (p-value < 0.05) improvements over the GPT and BART baselines, while a superscript of ** indicates statistically significant improvements over all baselines. We use Wilcoxon's signed rank test with a one sided alternative hypothesis to compute p values (Wilcoxon, 1992). baselines) generally perform worse than they do on the SBIC dataset. This is likely because the implicit hate corpus contains one reference per post, in contrast to the SBIC dataset (see Table 2). While the Expert knowledge model performs worse than the baseline, the other models perform slightly better. This is likely because the Expert model relies on toxicity classifications, which weren't available in the implicit hate corpus. We believe our models can be generalized to text generation tasks on other datasets, but they likely need multiple reference points where the implicit hate corpus only has one.
Explaining Toxic Text via Knowledge Enhanced Text Generation
We perform analyses and ablation studies on model results on the SBIC dataset. We do not perform these on the implicit hate dataset, since we have too few references per example. Examples of the error and challenge types below are given in Table  5. An additional full set of examples for each error and challenge type is given in Appendix 13. 
Explaining Toxic Text via Knowledge Enhanced Text Generation
We categorize the types of errors made by models on a small sample of 200 examples from the dev set and provide the distributions in Figure 3. We provide the error categories below.
Explaining Toxic Text via Knowledge Enhanced Text Generation
1. Non-Existent Stereotype: Model generates a stereotype when the reference stereotype is an empty string. The baseline GPT models tend to make more errors of every type, except that the expert model makes more errors of type 4. The expert model likely focuses on tokens that trigger toxicity classifications, which makes it less likely to focus on other relevant tokens. In the second example of Table 5, the token "black" may be triggering the Expert Model, causing an error of type 4.
Explaining Toxic Text via Knowledge Enhanced Text Generation
The knowledge enhanced models rarely fail to generate a stereotype (type 2 error). On the other hand, they often detect non existent stereotypes (type 1 error), but less often than the baselines. In the third example of Table 5, "contraceptive" may be incorrectly triggering the implicit and explicit models, while the expert model is not triggered.
Explaining Toxic Text via Knowledge Enhanced Text Generation
We categorize the various challenges faced by our text generation models and provide distributions in Figure 4 and counts in Appendix 15. We analyze the same sample of 200 examples from Section 5.1.
Explaining Toxic Text via Knowledge Enhanced Text Generation
The model fundamentally misunderstands the post and generates an irrelevant stereotype as a result.
Explaining Toxic Text via Knowledge Enhanced Text Generation
2. High Sensitivity: The model is highly sensitive to trigger words, which causes it to detect non-existent stereotypes or generate stereotypes based on the triggers alone.
Explaining Toxic Text via Knowledge Enhanced Text Generation
The model focuses only on parts of the input and generates stereotypes based on those parts, rather than on the entire input post.
Explaining Toxic Text via Knowledge Enhanced Text Generation
The model clearly considers the entire input post, but does not draw connections between the different parts of the post.
Explaining Toxic Text via Knowledge Enhanced Text Generation
The model takes a more literal interpretation of a sarcastic or ironic post, causing it to output text that has the opposite meaning of the reference stereotype.
Explaining Toxic Text via Knowledge Enhanced Text Generation
6. Ignores Stereotype: The model does not generate a stereotype despite a non-empty reference stereotype.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Interestingly, the MIXGEN model tends to misunderstand sarcasm and irony at a slightly higher rate than the other knowledge model types. In the fourth example of 5, MIXGEN and the Implicit Model take literal interpretations of the input post. The Implicit Model type has difficulty with drawing connections over the input (challenge type 4). An example is given in the 6th row of Table 5, where the model does not draw a connection between the target minority and the stereotypes present. Table 5 provides examples of MIXGEN synthesizing knowledge across sources. In the third example, MIXGEN produces an empty string even though two of the more reliable sources (explicit and implicit knowledge) produce non empty strings. In the fifth example, MIXGEN clearly combines parts of all the knowledge sources, while in the sixth example MIXGEN produces a more accurate stereotype than any of the other models. Thus, it seems that MIXGEN does not simply attempt to copy the correct model, but actually possess a deeper understanding of the knowledge types it synthesizes. Table 6 contains results of ablations on the Implicit Knowledge models. We vary the amount (k) and source (GPT and GPT-2) of implicit knowledge. The model which sources knowledge using GPT only outperforms the model sourcing knowledge using GPT-2 when k = 3. When k = 15, the latter model sourcing GPT-2 knowledge outperforms the model sourcing knowledge from GPT. It is possible that GPT-2 is less biased than GPT, thus benefiting our model for low k, but that these benefits decrease as k increases. With greater k, both models have a greater chance of exhibiting bias.    Table 6: BERTScores of baseline models and the implicit knowledge models.  In Table 7, we look at ablations on the number of knowledge informed models, for the MIXGEN CONCAT model. Let k be the number of knowledge informed models. The variant using k = 6 models performs best. Since the models variants within each knowledge type only vary by some hyperparameter, information gain probably saturates as the number of models increases; however performance decreases when k = 9. Since we tend to include models with lower performance for larger k, worse performing models are likely counter productive.
Explaining Toxic Text via Knowledge Enhanced Text Generation
We discuss limitations of our study here. Per the discussion in Subsection 4.3, we do not perform human evaluation of our results for multiple reasons. We do believe this is a limitation, and think future work may benefit from some form of human evaluation, while mitigating some of the concerns mentioned in Subsection 4.3.
Explaining Toxic Text via Knowledge Enhanced Text Generation
The MIXGEN model requires significant computational power. One needs to train models across knowledge sources, and then train the MIXGEN model itself. Future work may alleviate this burden by considering end to end solutions, or more efficient knowledge retrieval techniques.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Future work could perform an "in the wild" analysis, produced by procuring random comments from the internet and running our proposed models on these comments, to determine how effective the models might be in a real world setting. Further ablations may also provide insight. For instance, with respect to our implicit knowledge models, it may be helpful to remove BART altogether and instead use GPT for predicting the target minority, pretraining on implicit knowledge and for generating the final stereotype. Finally, it may be helpful to standardize the incorporation of knowledge into the models, such that the models using different knowledge types may be directly compared.
Explaining Toxic Text via Knowledge Enhanced Text Generation
In this paper, we propose a novel framework MIX-GEN to generate the stereotypes present in toxic social media posts, using multiple knowledge sources. We categorized three different sources of knowledge and synthesize the sources of knowledge using the MIXGEN models. While the knowledge models perform as well as baselines, models built on different knowledge types vary in strengths and weaknesses. For instance, the expert model suffers from high sensitivity to trigger words, while the implicit models may not draw connections over complex inputs. The MIXGEN models takes this into account and minimizes the number of examples on which it has errors and/or faces challenges. We conclude that mixture and ensemble methods as simple as concatenation can leverage the complementary nature of distinct knowledge sources to produce high quality text generations.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Models such as the one proposed in this paper, which output toxicity classifications of text or speech and reasoning behind such classifications should be used with care. Considerations of algorithmic fairness should be taken into account (Corbett-Davies et al., 2017), as well as cultural differences (Oliva et al., 2020) and racial biases (Xia et al., 2020) which can lead to misclassifications of offensiveness. Care should be taken to avoid political bias in training datasets, when training these models for deployment purposes (Wich et al., 2020). Finally, concerns about censorship should be taken seriously (Heins, 2014). also pass the input to the BART model and retrieve the BART encoded input, namely H ∈ R n×d . Let v 1 , . . . , v m in R d be trainable weight vectors. For the j-th row vector, h j of H, we compute an enriched hidden state h ′ j as follows:
Explaining Toxic Text via Knowledge Enhanced Text Generation
We then pass the enriched hidden state through the BART decoder to generate the output stereotype.
Explaining Toxic Text via Knowledge Enhanced Text Generation
To combine knowledge from multiple variables, we sum the enriched hidden state in Equation (2) over each variable.
Explaining Toxic Text via Knowledge Enhanced Text Generation
In this section, we provide formal details on how we incorporate explicit knowledge from Concept-NET using concatenation. In order to retrieve the top k triples (varying k ∈ {3, 5, 10, 15, 20, 25}), we first extract nouns, verbs, and adjectives from the input as our query tokens. We then query Con-ceptNet for triples associated with the query tokens and sort the triples by the product of the query's IDF weight and the triple's edge weight. We then translate the triples into English (Robyn Speer, 2019). For example, if the entities "car" and "vehicle" are connected by the edge relation "IsA", the translation would be "Car is a vehicle". We concatenate the translations to the input post to form a new input. Formally, let the input be "s post " and "s [trp_i] " be the sentence derived from the i-th triple. The modified input is then "s post [SEP]s [trp_1] [SEP] · · · [SEP]s [trp_k] ". We then pass the modified to the BART model which generates the implied stereotype.
Explaining Toxic Text via Knowledge Enhanced Text Generation
The concatenation based approach allows the model to encode the external knowledge into its own embedding space.
Explaining Toxic Text via Knowledge Enhanced Text Generation
In this section, we provide formal details on how we incorporate explicit knowledge from Concept-NET using attention and the fusion layer described in Chang et al. (2020). This is an alternative method we tried in addition to the method described in Section A.2.1. We first accumulate the top k triples (varying k ∈ {5, 10, 20}) associated with a post, using the method described in Appendix A.2.1. We then concatenate the numberbatch embeddings (each of dimension p) of the two entities in each of the k triples vertically to produce a vector in R 2p . We then horizontally stack the concatenations to produce a matrix, H KG ∈ R k×2p . The encoded input generated by the BART model can be represented by the matrix H B ∈ R n×d , where n is the input length and d is the hidden size of the BART model. We compute knowledge aware attention over the input as follows:
Explaining Toxic Text via Knowledge Enhanced Text Generation
Finally, we concatenate the original encoded input and the knowledge aware matrix, H B KG and perform an affine transformation:
Explaining Toxic Text via Knowledge Enhanced Text Generation
where ⊕ denotes column-wise concatenation, W 4 ∈ R 2d×d and b 4 ∈ R d . H ′ ∈ R n×d is the new hidden state, which fuses the old hidden state and the new knowledge aware hidden state. H ′ is passed into the BART decoder.
Explaining Toxic Text via Knowledge Enhanced Text Generation
A.3.1 Incorporating Implicit Knowledge using GPT
Explaining Toxic Text via Knowledge Enhanced Text Generation
In this section, we provide formal details on incorporating implicit knowledge using a pretrained GPT models. Let M tm be a BART model trained to predict the target minority of a given input post. Let the target minority string predicted by M [tm] be "s [tm] " and let "s [pr] " be a prompt. We provide a list of prompts in Table 8 We then prompt the GPT models with the following string "The s [tm] s [pr] ". The GPT models complete the prompt with a generated string, "s [gpt] ", so that the final string is "The s [tm] s [pr] s [gpt] ". For GPT generations, we choose hyper-parameters based on methods in Patrick von Platen (2020). For each input, we randomly select various prompts to generate k sentences of the form "The s [tm] s [pr] s [gpt] ", varying k ∈ {3, 15}. Note that if there is no predicted Prompt were known for were described as were regarded as were thought of as were well-known for were seen as a have a reputation for would frequently worked as have a job as have a part-time job as earned money by started working as have various hobbies such as would regularly engage in frequently talked about behaved as though liked to Table 8: Prompts for the GPT/GPT-2 biased sentence generators. The biased sentences are used to train the BART model, which is then retrained to produce the output stereotype. The predicted target minority of the input post is prepended to some of the prompts above, chosen at random. The GPT models then complete the prompt.
Explaining Toxic Text via Knowledge Enhanced Text Generation
target minority, s [tm] is the empty string and no sentence is generated using GPT. In this case, the input is paired with just the empty string.
Explaining Toxic Text via Knowledge Enhanced Text Generation
We then pair the input with each of the k generated sentences (or the empty string if there is no predicted target minority) and train a BART model M to predict the generated sentences. Model M is then used as a pretrained model and is retrained to predict the implied stereotype given the same input. The retrained BART model is our final implied stereotype generator.
Explaining Toxic Text via Knowledge Enhanced Text Generation
We train our BART models using a learning rate of 5e − 6 for 3 epochs with a batch size of 2 or 4, depending on the size of the input. This excludes the BERT classifier models, whose settings are given in Appendix B. The BART models have 406M parameters and all training is done on an Nvidia TITAN V GPU, with 12 GB memory, a boost clock speed of 1455 MHz, 640 Tensor cores and 5120 CUDA cores. Training under this regime takes approximately 90 -120 minutes. We remove all URLs, "RT" strings, and "@" mentions from the input post. We train these models on just a single seed and results are reported on just that seed, as we had limited time to train and test our models. The baseline GPT-2 and GPT models are trained for 5 epochs, as in the original paper by Sap et al. (2020). Following the paper, we perform minimal preprocessing to the input text before training and testing and only remove all URLs. During inference, we pass batches of input from the dev and test sets to the generate method of the huggingface BART model class. We use beam search for generation, with a beam width of 10 and a length penalty of 5.0.
Explaining Toxic Text via Knowledge Enhanced Text Generation
We train base BERT models, which have 110M parameters. Training takes approximately 30 -40 minutes on the GPUs described in B. We train these models on just a single seed, as results did not vary much as the seed varied and we had limited time to train and test our models.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Offensiveness 5e-6 32 2 Intent to Offend 5e-7 32 1 Lewdness 5e-6 32 1 Group Targeted 5e-6 32 2 BERT Classifier Model training settings are given in Table 9 and Table 10.
Explaining Toxic Text via Knowledge Enhanced Text Generation
In this section, we look at ablation studies conducted on the Expert Knowledge models and the Explicit Knowledge models.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Recall that we use the last attention layer of a BERT classifier with the join embedding architecture introduced in Pryzant et al. (2020) to enhance the hidden states of the BART model performing stereotype generation. We perform ablation studies by replacing the classifier over a few different annotated categories, namely Offensiveness, Intent to   Offend, Lewdness, and Group Targeted. We also train a model that uses all of the classifiers. The results for the ablations are in Table 11. It is important to note the relative performance of the models. The Expert Knowledge model using the Group Targeted BERT classifiers performs better than the other single classifier models, and performs on par with the Expert Knowledge model leveraging all of the classifiers. It's likely that the tokens used by the BERT model to identify whether a minority group is targeted aligns closely with the portions of the encoded input used to generate the stereotypes. This makes intuitive sense, since the set of posts targeting some minority group likely has some stereotype mentioned in the post and vice versa. The same cannot be said for the other categories. This intuition is further strengthened by the fact that the Expert Knowledge model using all the classifiers does not perform much better than the Expert Knowledge model using just the Group Targeted classifier. Thus the other classifiers may be contributing little additional knowledge to the stereotype generation task.
Explaining Toxic Text via Knowledge Enhanced Text Generation
In this section, we specifically discuss the number of knowledge triples we use when training ex-  plicit knowledge models and note trends in the BERTScore as k varies. The results are in Table  12. We discuss an additional attention based model not mentioned in the main paper. The detailed methodology for this model is given in Appendix A.2.2. We note that when concatenated directly to the input, performance increases as k increases up to a point and then starts to decrease. We believe that this occurs because the usefulness of knowledge initially increases then decreases as k increases. In particular, as k increases, many of the latter triples tend to be synonymous with earlier triples or unrelated to the original input. We also note that incorporating knowledge as attention tends to produce a better BERTScore, while not performing as well on other metrics. We suspect that this is due to the numerbatch embeddings capturing more semantic meaning than the BART embeddings. The numberbatch embeddings draw from a variety of sources, in addition to the graph structure itself and they perform well on benchmarks measuring word similarity.
Explaining Toxic Text via Knowledge Enhanced Text Generation
We provide a few examples of stereotype generation, comparing and contrasting the different model types in Table 13. jews were killed in the holocaust The token "jew" triggers an irrelevant generation.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Challenge Type 3 (Localized Generation) how do you know the asian food is undercooked? when it starts barking.
Explaining Toxic Text via Knowledge Enhanced Text Generation
asians have no food The post focuses on the first part of the sentence, but ignores the remainder.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Challenge Type 4 (Does not Draw Connections) enjoy having multiculturalism shoved down your throat, fighting for israel and having a black or asian as your new best friend.
Explaining Toxic Text via Knowledge Enhanced Text Generation
non-whites are inferior The generated stereotype understands that the stereotype is about non-whites, but does not connect it to the broader stereotype of non-white cultures having a negative impact.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Challenge Type 5 (Misunderstands Sarcasm/Irony) dating rules for women * don't be morbidly obese * women shouldn't be skinny The generated stereotype inverts the sarcasm of the original post and thus fails to capture the stereotype.
Explaining Toxic Text via Knowledge Enhanced Text Generation
Challenge Type 6 (Ignores Stereotype) fight anyone that tells you there's nothing wrong with being fat. and help your body and bones and lose weight, seriously.
Explaining Toxic Text via Knowledge Enhanced Text Generation
"" The input post targets those who are obese, yet the model does not generate a stereotype.  
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
Knowledge Graph (KG) with entities and relations provides rich knowledge for language learning (Wang et al., 2017(Wang et al., , 2014. Recently, researchers have explored to incorporate KG information into PLMs (Devlin et al., 2018;Radford et al.) to enhance language representations, such as ERNIE-THU (Zhang et al., 2019), WKLM (Xiong et al., 2019), KEPLER , KnowBERT (Peters et al., 2019), BERT-MK (He et al., 2019) and KALM (Rosset et al., 2020), . However, they only utilize entity information and ignore the finegrained relationships between entities. The finegrained semantic information of relations between entities is also critical to language representation learning. Taking Figure 1 as example, for entity typing, without explicitly knowing the fine-grained relation Guest between Lang Lang and Trio of Happiness, which is different from the relation Performer between Lang Lang and BBC Proms, it's impossible to correctly predict the type of Trio of Happiness as TV Show, since the input sentence literally implies that Trio of Happiness belongs to the same type as BBC Proms. The fine-grained relations between entities in KG provide specific constraint on entities, thus can play an important role in language learning for knowledge-driven tasks.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
To explicitly incorporate entities and finegrained relations in KG into PLMs, one main challenge we are faced with is the Text-Knowledge Alignment (TKA) problem: it's difficult to make token-relation and token-entity alignments for the fusion of text and knowledge. To handle this problem, the KG-enhanced pretrained language model (KLMo) is proposed to integrate KG (i.e. both entities and fine-grained relations) into the language representation learning. The main component of KLMo is a knowledge aggregator, which is responsible for text and knowledge information fusion from two individual embedding spaces, i.e. token embedding space and KG embedding space. The knowledge aggregator models the interaction between entity spans in text and all entities and relations in a contextual KG via an entity span-level cross-KG attention to make tokens attend to highly related entities and relations in KG. Based on the KG-enhanced token representations, a relation prediction objective is utilized to predict the relation of each pair of entities in text based on the distant supervision of KG. Furthermore, an entity linking objective is utilized to predict entities in KG based on the corresponding entity spans in text. The relation prediction and entity linking objectives are the key to the integration of KG information into text representations.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
We conduct experiments on two Chinese knowledge-driven NLP tasks, i.e. entity typing and relation classification. The experimental results demonstrate that KLMo obtains large improvements over BERT and existing knowledgeenhanced PLMs, by taking full advantage of a structured KG including both entities and fine-grained relations. We also will publish a Chinese entity typing dataset for the evaluation of Chinese PLMs.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
As shown in Figure 2, KLMo is designed as a multilayer Transformer-based (Vaswani et al., 2017) model, which accepts a token sequence and the entities and relations in its contextual KG as input. The token sequence is firstly encoded by a multi-layer Transformer-based text encoder. The output of the text encoder is further used as input for the knowledge aggregator that fuses the knowledge embeddings of entities and relations into the token sequence to obtain KG-enhanced token representations. Based on the KG-enhanced representations, novel relation prediction and entity linking objectives are jointly optimized as the pre-training objectives, which help incorporate high-related entity and relation information in the KG into the text representations.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
As shown in Figure 2, the knowledge aggregator is designed as an M -layer knowledge encoder to integrate knowledge in KG into language representation learning. It accepts the hidden embeddings of the token sequence and the knowledge embeddings of the entities and relations in KG as input, and fuses text and KG information from two individual embedding spaces. The knowledge aggregator contains two separate multi-head attentions: token-level self-attention and knowledge graph attention (Veličković et al., 2017), which encodes the input text and the KG independently. The entity
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
Lost and Love is a movie starring Andy Lau (playing Lei Zekuan) e1 e2 e3 
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
invisible Entity Spans
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
(2) Entity Linking MLP
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
(1) Relation Prediction MLP Actor 33165 Figure 2: Overview of the model architecture.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
representation is computed by pooling over all tokens in an entity-span. Then the aggregator models the interaction between entity spans in text and all entities and relations in a contextual KG through an entity-level cross-KG attention to incorporate knowledge into the text representations.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
Knowledge Graph Attention As the entities and relations in a KG composes a graph, it's critical to considering the graph structure during knowledge representation learning. We first represent entities and relations in the contextual KG by TransE (Bordes et al., 2013) and then translate them into an entity and relation embedding sequence {z 0 , z 1 , ..., z q }, served as the input for the knowledge aggregator. Then the knowledge aggregator encodes the entity and relation sequence by a knowledge graph attention which considers its graph structure by importing a visible matrix M into the traditional self-attention mechanism (Liu et al., 2020). The visible matrix M only allows adjacent entities and relations in the KG to be visible to each other during representation learning, as shown in the right bottom of Figure 2.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
Entity-level Cross-KG Attention To compute the KG-enhanced entity representations, given an entity mention list C e = {(e 0 , start 0 , end 0 ), ..., (e m , start m , end m )}, the knowledge aggregator first computes the entity span representations {ê i 0 , ...,ê i m } by pooling over all tokens in an entity-span with self-attentive span pooling method from (Lee et al., 2017). The entity span embeddings {ê i 0 , ...,ê i m } can be expanded to all tokens {ê i 0 , ...,ê i n } by makingê i j =t i j for tokens not in any entity spans, wheret i j denotes the representation of the j-th token from the token-level self-attention.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
In order to model the interaction between entity spans in text and all entities and relations in a contextual KG, the aggregator performs an entity-level cross-KG attention to allow tokens attend to highly related entities and relations in KG, thus computes the KG-enhanced entity representations. Specifically, the entity-level cross-KG attention in the i-th aggregator is performed by contextual multihead attention between the entity span embeddings {ê i 0 , ...,ê i n } as the query and the entity and relation embeddings {z i 0 , ..., z i q } as the key and value. KG-enhanced Token Representations To inject the KG-enhanced entity information into the token representations, the i-th layer of the knowledge aggregator computes the KG-enhanced token representations {t i 0 , ...,t i n } by adopting an information fusion operation between {t i 0 , ...,t i n } and {e i 0 , ..., e i n }. For the j-th token, the fusion operation is defined as follows:
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
where u i j represents the hidden state integrating the information from both token and entity. σ is a nonlinear activation function. W i * and b i * are learnable weights and biases respectively. The KG-enhanced token representation {t i 0 , ...,t i n } is fed into the next layer of knowledge aggregator as input.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
To incorporate KG knowledge into the language representation learning, KLMo adopts a multi-task loss function as the training objective:
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
In addition to the loss of masked language model L M LM (Devlin et al., 2018;, an relation prediction loss L RP and an entity linking loss L EL are integrated to predict the entities in KG based on the corresponding KG-enhanced tokens representations {t M 0 , ...,t M n }. For each pair of entity spans, we utilize the relation between their corresponding entities in the KG as the distant supervision for relation prediction. The relation prediction and entity linking objectives are the key to the integration of relations and entities in KG into the text. Since the number of entities in KG is quite large for the Sof tmax operation in entity-linking objective, we 
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
This section presents the details of KLMo pre-training and its finetuning on two specific knowledge-driven NLP tasks: entity typing and relation classification. We pretrain KLMo by a Chinese corpus of Baidu Baike's webpages and the Baike Knowledge Graph. Details of the pretraining corpus and experimental settings are described in Appendix A. 1
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
We compare KLMo with the state-of-the-art PLMs pretrained on the same Baidu Baike corpus: (1) BERT-Base Chinese (Devlin et al., 2018), which is further pretrained on the Baidu Baike corpus for one epoch. (2)    (1) All knowledge-enhanced PLMs generally perform much better than the BERT baseline on all measures, which shows that entity knowledge is beneficial to entity type predication with limited annotated resources.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
(2) Compared with the existing knowledge-enhanced PLMs, KLMo largely improves the recall score over WKLM and ERNIE, leading to an improvement of 1.58 and 0.57 on micro-F1 respectively. This indicates that finegrained relationships between entities help KLMo to predict appropriate categories for more entities.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
Dataset The CCKS 2019 Task 3 Inter-Personal Relational Extraction (IPRE) dataset (Han et al., 2020) is used for the evaluation on relation classification. The training set is automatically labeled by distant supervision, and the test set is manually annotated. There are 35 relations (including a null-relation class "NA"), where "NA" accounts for nearly 86% in the training set and 97% in the test set. The detail statistics of the dataset and finetuning settings are shown in Appendix B.2.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
We adopt precision, recall and micro-F1 as the evaluation measures. The results are shown in Table 2. In addition to BERT baseline, we also compare KLMo with an official CNN baseline, which gets CNN output as the sentence embedding and feed it into a relation classifier. From Table   2, we can see that both CNN and BERT baseline models do not perform well, which indicates the high difficulty of the dataset. This ascribes to the large number of noisy labels in the training set automatically generated by distant supervision. Although the dataset are very difficult, we can still observe that: (1) All knowledge-enhanced PLMs largely improve the precision and micro-F1 scores over BERT baseline, which shows that both entity information and KG information can enhance language representations and accordingly prompt the performance of relation classification.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
(2) KLMo largely improves the precision score over WKLM and ERNIE, leading to an improvement of 2.41 and 1.29 on micro-F1 respectively, which demonstrates that fine-grained relations in KG help KLMo avoid fitting on noisy labels and predict relations correctly.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
Most NLP tasks only provide text inputs and the entity linking itself is a hard task. Thus, we investigate the effects of KG entities and relations for KLMo on entity typing. w/o KG refers to finetuning KLMo without the input of KG entities and relations. Table 3 shows the results of the ablation study. Without KG input for funetuning, KLMo still largely outperforms BERT on both precision and recall scores, leading to an improvement by 1.74 on micro-F1. Compared with KLMo finetuning with KG, KLMo without KG witnesses a small decrease of 0.84 on micro-F1 measure. This demonstrates that KG information has been integrated into KLMo during pre-training. For most specific NLP tasks, KLMo can be finetuned in a similar way as BERT.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
In this paper, we propose a novel KG-enhanced pretrained language model KLMo to explicitly integrate KG entities and fine-grained relations into the language representation learning. Accordingly, the novel knowledge aggregator is designed to handle the heterogeneous information fusion and textknowledge alignment problems. Further, the relation prediction and entity linking objectives are jointly optimized to encourage the knowledge information integration. The experiment results show that KLMo outperforms the other state-of-the-art knowledge-enhanced PLMs, which validates the intuition that fine-grained relationships in KG can enhance the language representation learning and benefit some knowledge-driven NLP tasks.   
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
In the experiment, we first obtain the knowledge representations trained on Baike KG triples by TransE (Bordes et al., 2013) algorithm using the OpenKE toolkit (Han et al., 2018). The representations are used to initialize the entity and relation embeddings in KLMo. The embedding dimension is set to 100 and the epoch number is set to 5000. As for the pre-training of KLMo, due to the expensive cost of pre-training from scratch, we inherit the parameters of BERT-Base Chinese to initialize the Transformer blocks for token encoding, while the parameters for entity and relation encoding modules are all randomly initialized. The number of text encoder layers L and knowledge aggregator layers M are both 6. The hidden size of token embeddings d t , knowledge embeddings d z and entity span embeddings d e are set to 768, 100 and 100. The number of token-oriented attention heads A t , KG-oriented attention heads A z and entity span-level attention heads A e are set to 12, 4 and 12 respectively. The pre-training of KLMo runs 3 epochs on 4 NVIDIA Tesla V100 (32GB) GPUs with the batch size of 128, the max sequence length of 512 and the learning rate of 5e-5.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
To evaluate the performance of KLMo, two knowledge-driven tasks, i.e. entity typing and relation classification, are performed in this work.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
Given a sentence with an entity mention, the entity typing task is to label the mention with its finegrained semantic type.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
Dataset Entity typing is not a new task. However, to our best knowledge, there is no public benchmark dataset available on Chinese fine-grained entity typing. Therefore, In this work, we create a Chinese entity typing dataset, which is a completely manually-annotated dataset containing 23,100 sentences and 28,093 annotated entities distributed in 15 fine-grained categories of media works, such as Movie, Show and TV Play. We split the dataset into a training set with 15,000 sentences and a test set with 8,100 sentences. The detail statistics of the dataset are shown in Table 4.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
Finetuning The Chinese entity typing dataset lacks of KG entity annotations, thus we first use an entity linker tool accompanied with Baike Knowledge Base to recognize entity mentions in sentences and link them to their corresponding Baike KG entities. The statistics of the linked entity typing dataset are shown in Table 4. Over 50% of sentences contain at least one linked KG entity mention in both training set and test set. To finetune KLMo for entity typing, we use the representation of the first token of each entity span to predict its entity type. The model is finetuned for 10 epochs on the training set with the batch size of 128, the max sequence length of 256 and the learning rate of 2e-5.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
We also compare the results of various pretrained models on the task of relation classification. Given a pair of entities in a sentence, the relation classification task is to determine the relation type between the pair of entities.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
Dataset The CCKS 2019 Task 3 Inter-Personal Relational Extraction (IPRE) dataset (Han et al., 2020) is used for the evaluation on relation classification. The training set is automatically labeled by distant supervision, and the test set is manually annotated. There are 35 relations (including a null-relation class "NA"), where "NA" accounts for nearly 86% in the training set and 97% in the test set. The detail statistics of the dataset are shown in Table 5.
KLMo: Knowledge Graph Enhanced Pretrained Language Model with Fine-Grained Relationships
Finetuning The IPRE dataset also lacks of KG entity annotations, and we recognize and link entity mentions to their corresponding Baike KG entities in the same way as we do for the entity typing dataset. The statistics of the linked dataset are shown in Table 5. Over 40% of sentences contain at least one linked KG entity mention. To finetune KLMo for relation classification, we concatenate the representations of the first token of the two candidate entity spans. The model is finetuned for 10 epochs with the batch size of 128, the max sequence length of 256 and the learning rate of 2e-5.