Pretrain Knowledge-Aware Language Models

1 INTRODUCTION . The strong effectiveness and rich generalization ability of pretrained language models ( PLMs ) ( 1 ; 2 ; 3 ; 4 ; 5 ) have raised many questions about what is captured in transformer networks and why . Recent explorations found the pretrained language models may “ rediscover ” the linguistic pipeline at various transformer layers ( 6 ) , can serve as implicit knowledge bases for relation extraction ( 7 ) , perform soft reasoning tasks ( 8 ; 9 ) , and conduct some language tasks reasonably in a fully unsupervised , zero-shot , fashion ( 3 ; 10 ) . With sufficiently large amount of parameters , i.e . several billions , and enough task-specific supervision , the pretrained language models can even directly generate answers for natural language questions , at the same accuracy with state-of-the-art reading comprehension systems , without using any context documents or knowledge graphs ( 11 ) . Impressive as they are , language models are still far from ready to serve as an “ unsupervised multi-task learner ” that learns knowledge directly from human language and generalizes to downstream language tasks ( 3 ) . There are notable gaps of language models ’ performances on downstream tasks between models with ( 11 ; 12 ) and without ( 3 ) large amounts of task-specific fine-tuning . The language models still ( de- ) generate dull , factually-incorrect , or dream-like text when used in natural language generation ( 13 ; 14 ; 15 ; 16 ) . These challenges often necessitate over-parameterization ( 17 ) , grounding on external structural semantics ( 14 ; 16 ; 18 ; 19 ) , or large amount of task-specific fine-tuning ( 11 ) , which are costly , complicated , and not always feasible for every language task . One potential limitation of these language models is their style of pretraining , e.g , auto-regressive language modeling ( 3 ) or masked language modeling ( 1 ) , wherein transformer networks process a sequence of words and are asked to predict the next/masked words . There is no explicit guidance to the transformers that humans prefer them to capture correct , real-world information . As a result , all the knowledge captured in these pretrained language models is only signaled by patterns of co-occuring words in the input sequence that is learned implicitly during pretraining . In this paper , instead of creating bigger models or adding knowledge-specific architectures , we propose to more efficiently leverage the existing parameters in the standard transformer language model , by simply making them aware of the various forms an entity can manifest itself as , and its role in the surrounding text . More specifically , this knowledge-awareness is communicated via the input fed to PLMs and in the output expected from them during pretraining . For input-awareness , we use an entity-name ( surface form ) dictionary that tokenizes word spans to their most popularly referred-to entity , e.g. , as fuzzy frequency-based entity annotations ( 20 ) , and serve these entity tokens as a parallel input channel along with the word tokens . For output-awareness , in addition to the language modeling objective , we add an entity prediction task that guides the model to distinguish the correct entity from various negative distractions . The two objectives together explicitly guide the language model to predict not only the correct words , but also the correct entity behind those words during pretraining , without changing the network architecture . By adding knowledge awareness to GPT-2 style auto-regressive language models , our pretrained language model , “ Knowledge-Aware Language Model ” ( KALM ) , shows significantly improved handling of knowledge-sensitive tasks . In the LAMA knowledge probing tasks ( 7 ) , KALM outperforms its entity-unaware baseline , GPT-2 , by about 25 % across all tasks at both base and large transformer sizes . Our 24 layer KALM ( Large ) is even comparable with the 17 Billion parameter GPT-2 on some tasks . It more accurately captures commonsense knowledge , factual semantics , and also relation semantics in these LAMA tests . The knowledge signals also aid generic language understanding : we have observed better language modeling perplexity and word prediction accuracy with KALM too . The advantages in language modeling also transfer to downstream tasks . In zero-shot question answering , the exact match accuracy of the answers generated by KALM are 20 % -100 % better than those of an equivalent GPT-2 model . We did not use any task-specific supervision or additional gradient updates , relying solely on the unsupervised knowledge learned in KALM . We only feed in a few example question-answer pairs as templates to format how generated answers should look . Injecting rich knowledge signals leads to improvements approximately equal to those gained by doubling the transformer layers , indicating that PLMs can be trained more efficiently – growing the parameters exponentially is not the only way to improve language understanding . To better understand pretraining and the advantage of knowledge awareness , we leverage the edge probe technique ( 6 ; 21 ) and dissect what is learned in the representations at various gradient step numbers throughout pretraining . We observe that the auto-regressive transformers start to learn the basis of language in the beginning of the pretraining , and gradually learns more complex semantics in the process ; adding knowledge-awareness greatly accelerates learning of higher level semantics , e.g. , coreferences and entity types , and helps the model perform better in those more complicated tasks . 2 PRETRAINING KNOWLEDGE-AWARE LANGUAGE MODELS . In this section we first present preliminaries in language modeling and then how we add knowledgeawareness in their pretraining . 2.1 PRELIMINARY . In this paper , without loss of generality , we mainly focus on the auto-regressive language modeling . Considering the text X as a sequence of tokens ( words or sub-words ) : X = { w1 , ... , wi , ... , wn } , the classical unidirectional factorization of language probabilities ( 22 ; 23 ) describes : p ( X ) = ∏ i p ( wi|w < i ) , ( 1 ) where w < i refers to all the tokens appear before i . This conditional probability can be parameterized in various ways . An effective choice is to use the uni-directional transformer , as done in GPT-2 ( 3 ) : p ( wi|w < i ) = transformer ( wi|w < i ) . The language modeling task provides a large amount of data to pretrain very deep transformer networks ( 5 ; 24 ; 25 ) . Scaling up the transformer parameter sizes will lead to significant improvements in the language model capability : with wider and deeper transformer layers , it is observed that transformer language models start to output more complicated semantics beyond lexical and syntactic patterns ( 6 ; 8 ; 7 ) . On the other hand , a roughly log-linear relationship between transformer size and output quality has been established , e.g . doubling the quality requires ten times more parameters and training data ( 3 ; 17 ; 10 ) . Even in industry , the marginal gain of increasing parameters will eventually be outweighed by the cost to train and serve such models . 2.2 KNOWLEDGE-AWARE PRETRAINING . As shown in Eqn . 1 , the pretraining is solely at the token level . All the semantics in PLMs are captured by the transformer indirectly ; there is no explicit requirement in pretraining to better capture knowledge – yet we expect them capture knowledge beneath the raw word sequences implicitly , e.g. , to generate factually correct statements . In this work we mitigate this discrepancy by making transformer networks aware of knowledge in language model pretraining . Instead of stacking more layers or adding external knowledge storage , we present a knowledge-aware language modeling ( KALM ) framework that packs more information into the same amount of transformer parameters . The first step to introduce knowledge awareness is an entity tokenizer that forms an additional entity token sequence ( 26 ) to signal the existence of entities in the input and output of the pretraining process . Entity Tokenizer . An entity tokenizer segments the text sequence into entity ids using a surface form dictionary , which maps word-ngrams to entities : wi : i+k dict look up−−−−−−→ ei , where ei is the most popular entity referred by the word k-gram wi : i+k , and ei = null if wi is not part of any known entity surface names . This simple dictionary look-up can be conducted efficiently , similar to the ( sub ) word tokenizer . Simultaneously , the text is tokenized into two channels – a word-entity duet token sequence ( 26 ) : Xduet = { { w1 , ... , wi , ... , wT } Word Sequence ; { e1 , ... , ei , ... , eT } Entity Sequence . ( 2 ) The two sequences are aligned position by position . If multiple ( sub ) words together form an entity name , the corresponding entity id is duplicated in each position corresponding to these words . For example , the name “ United States ” at wi : i+2 is mapped to entity “ USA ” in ei and ei+1 . Instead of enlisting a more precise entity linker or supervisions from entity labels ( 16 ; 18 ) , a fuzzy frequency-based dictionary look up places higher expectations on the model to use clues in the text to jointly build token and entity representations that better reflect how language conveys knowledge . Using a highly tuned entity linker would propagate its own biases into the transformer . Knowledge-Aware Input . Just as there is an input embedding for every word token , we allow the model to learn an entity embedding for each entity : ~ei = Embeddinge ( ei ) ∈ Rde , ( 3 ) ~wi = Embeddingw ( wi ) ∈ Rdw . ( 4 ) The two embeddings are combined to form the knowledge aware input : ~ti = ~wi + Lineart ( ~ei ) , Lineart ∈ Rde×dw . ( 5 ) All the embeddings are randomly initialized and learned in pretraining . Knowledge-Aware Output . The knowledge-aware input is fed into standard transformer layers , the same as the word-only input . Then in pretraining , besides the next-word prediction task , we also employ a next-entity prediction task to further incorporate knowledge-awareness . Specifically , we use one output head for the word probability , one for the entity , and share all transformer layers between words and entities . If there are L transformer layers and hLi is the output of the final layer ’ s i-th token , the loss for position i is computed as le ( ei|t < i ) = max ( 0 , s ( ~hi L , ~ei ) − s ( ~hi L , ~e− ) + λ ) , ( 6 ) s ( ~hi L , ~ej ) = cos ( Linear ( ~hi L ) , ~ej ) , ( 7 ) ~hi L = transformerL ( t < i ) . ( 8 ) The transformers in Eqn . 8 are stacked multiple times similar to GPT-2 . The scoring function in Eqn . 7 projects the hidden state into the embedding space with a Linear layer and takes the cosine similarity with an arbitrary entity ej . Assuming that position i is linked by our entity linker to ei , we carefully choose a corresponding negative entity e− to contrast with ei using margin loss with margin λ in Eqn . 6 . Pretraining . The knowledge-aware input and output are incorporated in the standard multi-task set up for our KALM knowledge-aware language model pretraining : lKALM ( Xduet ) = ∑ i lw ( p ( wi|t < i ) ) + αle ( ei|t < i ) . ( 9 ) It combines the language modeling loss lw ( ) – cross-entropy as in standard PLMs – with the entity prediction loss le ( ) , where α is a hyper-parameter to balance the two . Inference . At inference time – whenever generating output text – KALM use the word prediction head p ( wi|t < i ) , which is consistent with GPT-2 . The entity prediction task is an auxiliary task that guides the model to attend to the entity semantics during pretraining ; in inference only the shared transformer representations is used upon the input word and entity tokens . Compared with the standard GPT-2 style PLMs , the architecture of KALM only differs with an enlarged tokenization vocabulary with additional entity tokens , and their entity embeddings before the input to the transformer network . The transformer architecture and its layer configurations are kept consistent .

Summary: This paper proposes KALM - a knowledge-aware language model that incorporates entity information. Specifically, the method involves using a dictionary lookup to match n-grams to entities from a database. Then, each token is embedded into two vectors — one using the surface form, and one using the entity ID it is matched to (if there is a match). These embeddings are then added and passed into a Transformer model to predict the next word given the context. Experiments compare the model to GPT-2 and T5 (two state-of-the art Transformer-based LMs) and demonstrate the benefit of adding entity information.