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 presents a knowledge-aware language model pretraining method without changing model architecture. Specifically, they add entity prediction task along with language modeling task to make the model aware of knowledge. Experiments show improved results on the LAMA knowledge probing task compared to GPT-2 models. They also show comparable results on zero-shot question answering task, even with only 2% transformer parameters compared to GPT-2 17B.