Knowledge Infused Decoding

1 INTRODUCTION . Pre-trained language models ( LMs ) have been shown to capture rich semantic and syntactic features , as demonstrated by the state-of-the-art performance on many generation tasks such as abstractive summarization ( Zhang et al. , 2020b ; Chu & Liu , 2019 ) and dialogue generation ( Roller et al. , 2021 ; Zhang et al. , 2020c ) . Their generations , however , can be quite limited in scenarios requiring knowledge—they frequently struggle with issues such as being easily misguided by phrase cooccurrence ( e.g. , “ Talk ? Birds can talk. ” ) , failing to handle negation ( e.g. , “ The theory of relativity was not developed by Einstein. ” ) ( Kassner & Schütze , 2020 ) , and being unable to compare commonsense concepts , such as time ( Qin et al. , 2021 ) and digits ( Talmor et al. , 2020 ) . To enhance the performance of LMs on knowledge-intensive NLG tasks1 , prior studies have proposed to re-train LMs with knowledge-aware objectives ( Zhou et al. , 2020 ; Xiong et al. , 2020 ; Zhang et al. , 2019 ; Khandelwal et al. , 2020 ) or add special architectures to encode knowledge ( Bosselut et al. , 2019 ; Logan et al. , 2019 ; Peters et al. , 2019b ) from external resources ( e.g. , knowledge graphs such as CONCEPTNET ( Speer et al. , 2017 ) and ATOMIC ( Sap et al. , 2019 ) ) . These methods , though yielding impressive results on many downstream tasks , can be computationally expensive . More importantly , knowledge implicitly parameterized in LM architectures is difficult to revise and expand ( Lewis et al. , 2020b ) , and wrong generations are hard to diagnose due to lack of interpretation ( Talmor et al. , 2020 ) , which heavily limits their real-world applications . More recent retrieval-based models try to tackle these problems by augmenting inputs with retrieved knowledge evidence ( Lee et al. , 2019 ; Guu et al. , 2020 ) . For example , RAG ( Lewis et al. , 2020b ) leverages non-parametric memory to access extensive knowledge ( in the form of unstructured documents ) , 1We define knowledge-intensive NLG tasks as those whose input context alone does not provide complete knowledge for a legitimate and plausible generation . and jointly fine-tunes a parametric LM ( i.e. , BART ( Lewis et al. , 2020a ) ) to enable knowledge-aware generation . A key limitation of such methods is that they retrieve documents only once while grounding them in the input static context , and thus can not support the dynamic nature of the context as new tokens are generated . The static knowledge becomes a major problem for tasks where longer and abstractive generation is expected , such as open-ended story generation ( Mostafazadeh et al. , 2016 ) , multi-turn dialogues ( Zhao et al. , 2020 ) , and conversation summarization ( Gliwa et al. , 2019 ) . Moreover , in a recent study , Krishna et al . ( 2021 ) replaced the knowledge retriever in RAG with a random retriever and found little difference in the resulting performance on a long-form QA task named ELI5 ( Fan et al. , 2019b ) , indicating the model may not be actually grounding its text generation to the retrieved documents . To address these limitations , in this work , we present a novel decoding algorithm KID , aiming to better infuse knowledge into generation in a dynamic manner . Instead of solely relying on the static knowledge retrieved at beginning , during each step of LM decoding , KID dynamically searches promising continuation from retrieved knowledge , to guide the current step generation . Specifically , KID maintains a local knowledge memory , interacts it with a knowledge trie dynamically created from retrieved supporting documents , and updates the local memory as a knowledge-aware constraint to guide the generation . The key intuition behind KID is that existing LM pre-training objectives are usually defined at the token level yet do not explicitly model concept-centric knowledge ( Xiong et al. , 2020 ) — thus motivating us to reshape the probability mass at each step decoding towards the distribution of entities in knowledge . The contributions of this work are three-fold : First , we introduce KID as a model and task agnostic decoding method that integrates knowledge on the fly and can be applied to various knowledgeintensive tasks with different generative LMs . Second , from a docoding perspective , on six knowledgeintensive NLG tasks , GPT2 ( Radford et al. , 2019 ) and BART ( Lewis et al. , 2020a ) equipped with KID significantly outperform conventional beam search or sampling decoding by a large margin . Third , from a knowledge infusion perspective , unlike six strong knowledge-infusion baselines which require either additional retraining or special architecture modifications , KID leverages knowledge more effectively as a light-weight knowledge infusion solution . Additionally , in few-shot scenarios KID significantly improves over them , demonstrating its generalization ability in low-resource and domain shifting regimes . 2 RELATED WORK . We briefly review existing work enhancing LMs with external knowledge and representative decoding algorithms for generation . Enhancing Language Model with Knowledge . Large language models implicitly encode knowledge in their parameters but with limits ( Petroni et al. , 2019 ; Kassner & Schütze , 2020 ; Lin et al. , 2020a ) . Several architectures and objective functions have been proposed to explicitly encode external knowledge ( Sun et al. , 2020 ; Logan et al. , 2019 ; Roberts et al. , 2020 ; Levine et al. , 2020 ) , or to augment LM pre-training data with retrieved knowledge ( Lewis et al. , 2020b ; Guu et al. , 2020 ; Lee et al. , 2019 ) . However , Talmor et al . ( 2020 ) notes that the reasoning ability of such LMs is strongly tied to the context seen during pre-training and is thus hard to generalize to new domains . Built on LMs , KID incorporates extra knowledge from external resources ( Wikipedia ) and thus shows strong performance in knowledge-intensive NLG tasks . Better Decoding Algorithm . Two common strategies dominate the decoding algorithms used by most generative models : beam search which maximizes likelihood in a local horizon ( due to finite beam size ) , and sampling decoding ( e.g. , top-k sampling ( Fan et al. , 2018 ; Holtzman et al. , 2018 ) ) which relies on randomness . Holtzman et al . ( 2020 ) find beam search often produces generic and repetitive generation , while top-k sampling tends to pick unlikely tokens which creates incoherent and unrelated sequences . Existing attempts to mitigate these problems include reranking ( Adiwardana et al. , 2020 ; Shao et al. , 2017 ) , adding control signals ( Zhang et al. , 2018 ; Xing et al. , 2017 ) , and self-adaptive truncation ( Welleck et al. , 2020 ; Peters et al. , 2019a ) . None of these decoding algorithms consider integrating knowledge in the generation process . Reflective decoding ( West et al. , 2021 ) and DeLorean ( Qin et al. , 2020 ) are two recent decoding algorithms that focus on abductive commonsense reasoning . Reflective decoding in particular has the potential to be extended to other knowledge-intensive tasks . We compare it with KID in our experiments . 3 KNOWLEDGE INFUSED DECODING . We detail the implementation of KID in this section . As shown in Figure 1 , KID comprises of three steps : retrieving relevant knowledge ( §3.1 ) , constructing external and local knowledge memory ( §3.2 ) , and guiding current step decoding under the constraint of the knowledge trie ( §3.3 ) . 3.1 RETRIEVING EXTENSIVE KNOWLEDGE FROM WIKIPEDIA . The first step of KID is to retrieve several context-relevant documents to ground the following generation . We use DPR ( Karpukhin et al. , 2020 ) as our general-purpose knowledge retriever , which projects contexts and relevant documents to a 768-dim shared embedding space using a bi-encoder network ( i.e. , two independent BERTs ( Devlin et al. , 2019 ) ) . Here , the documents refer to the 100- word chunks of Wikipedia passages released by RAG ( Lewis et al. , 2020b ) , a total of 21M documents as our knowledge source Z . We pre-load the weights from the latest checkpoint of DPR ( March 2021 ) , as it improves retrieval performance by using mined negative samples and contrastive learning , which is also suggested by Jernite ( 2020 ) . During retrieval , we perform a maximum inner-product search ( MIPS ) with faiss2 accelerated by GPU ( Johnson et al. , 2019 ) . Formally , we retrieve k most relevant document z [ 1 , ... , k ] ∈ Z for context xcontext as : z [ 1 , ... , k ] = { zi ∈ Z|topk { BERT ( xcontext ) ⊤ · BERT ( zi ) } } ( 1 ) where BERT ( · ) means the vectors encoded by BERT . The number of retrieved documents k is a task-specific hyper-parameter—we discuss its impact on performance in §4.3 . 3.2 CONSTRUCTING EXTERNAL AND LOCAL KNOWLEDGE MEMORIES . We convert multi-document input retrieved from previous step into compressed knowledge memories , in order to 1 ) allow relevant knowledge to be easily identified , 2 ) reduce the memory footprint of the knowledge , whose size grows linearly with the number of retrieved documents . Following design choice of previous successful methods ( Bosselut et al. , 2021 ; Huang et al. , 2020 ; Fan et al. , 2019a ) , we adopt co-reference resolution and open information extraction ( OpenIE ) ( Stanovsky et al. , 2018 ) to convert plain text into triplet form3 . For example , knowledge statement like “ Iceland is a Nordic island country in the North Atlantic Ocean and it is the most sparsely populated country in Europe. ” will be converted to subject-relation-object triplets such as ⟨subj : Iceland , rel : is , obj : Nordic island country⟩ , ⟨subj : Iceland , rel : is , obj : most sparsely populated country in Europe⟩ , etc . To account for overlapping elements from these triplets , we use a prefix tree ( namely Knowledge Trie Gext ) to store and organize the extracted triplets , by setting the subject in each triplet as the key in the Gext . Note that unlike common character-level dictionary tries ( Chen et al. , 2020 ; Germann et al. , 2009 ) , in Gext each triplet is stored in token unit as our goal is to efficiently query and traverse the knowledge triplets stored in it . A tree structure encoding knowledge is appealing to knowledge intensive NLG tasks , since 1 ) the non-cyclic structure helps reduce repetitions in generations , and 2 ) querying a prefix tree can be efficiently completed in constant time ( O ( |xcontext| ) ) which does not involve any costly traversal on the graph ( Ji et al. , 2020 ; Zhang et al. , 2020a ) , regardless of the scale of grounding knowledge ( normally |xcontext| ≪ |k Wiki docs| ) . We also maintain a local memory Gloc ( a first-in-first-out list ) that records all the mentioned entities in current context , to focus on concept-centric knowledge ( Zhou et al. , 2020 ; Lin et al. , 2020b ) . More details about how we construct and query these knowledge memories can be found in §A.2 of Appendix .

Summary: This paper deals with the well-known but important problem of language models -- enhancing language model with (external) knowledge. To that end, the authors present Knowledge Infused Decoding (KID), a novel decoding algorithm for generative LMs, which infuses external knowledge into each step of the LM decoding. This method maintains a local knowledge memory based on the current context and continuously updates the local memory via reinforcement learning to guide each step of language generation. As a result, it can integrate knowledge on the fly and perform well on various knowledge-intensive tasks such as abstractive QA, logic-centric writing, and dialogue generation.