Mention Memory: incorporating textual knowledge into Transformers through entity mention attention

1 INTRODUCTION . 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 ( Cao et al. , 2021 ) . 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 ; Sun et al. , 2021 ) . 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 ) . 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 corpus . This retrieval is much more fine-grained than aggregate entity retrieval methods such as Entities as Experts ( EaE ) ( Fé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 ques- tion 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 . 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 ) and FM2 ( 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 . 2 METHOD . 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 . 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 . Notation . An input to the model is a passage x = x1 , . . . , xT 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 [ Estart ] and [ Eend ] 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 [ Estart ] nationality [ Eend ] of the [ Estart ] hero [ Eend ] who killed [ Estart ] Medusa [ Eend ] ” . Each mentionm 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 xs = [ Estart ] and xe = [ Eend ] . Representations of these tokens are later used to generate mention encodings . 2.1 CONSTRUCTING MENTION MEMORY FROM CORPUS . 2.1.1 MENTION ENCODER . Let H ∈ RT×d be token representations where d is the hidden dimension , such that Hi ∈ Rd is the contextualized embedding for the i-th token . Following Févry et al . ( 2020 ) we compute the encoding of a span ( s , e ) as a learnable linear projection W of the concatenation of its start and end token representations Hs and He SpanEncodingLayer ( H , ( s , e ) ) =W [ Hs ; He ] ( 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 . 2.1.2 MENTION MEMORY . 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 ∈ RN×dK and MemValue ∈ RN×dV , where N is the total number of linked entity mentions in English Wikipedia ( approximately 150 million ) and dK and dV are dimensions of key and value encodings . Additionally , we record entity ( Wikipedia ) IDs of mentions in MemEnt ∈ RN , 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 . 2.2 TOME MODEL . 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 ) , ( 2 ) H ′ = TransformerBlock ( M ) ( 3 ) The TOME model encodes input passages x with the word embedding layer and initial Transformer block and then applies one or more TOMEBlocks H0 = InitialTransformerBlock ( TokenEmbedding ( x ) ) , ( 4 ) H l = TOMEBlockl ( H l−1 ) , l = 1 . . . L ( 5 ) 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 . 2.2.1 MemoryAttention 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 representationM ( omitting layer indices ) . Consider a mentionm 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 ) ) , ( 6 ) Query mention encodings are used to retrieve relevant memories from the Mention Memory table . 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 αi ∝ exp ( Query ( m ) · MemKey ( i ) ) , i ∈ TopMem ( Query ( m ) ) ( 7 ) Value ( m ) = ∑ i∈TopMem ( Query ( m ) ) αi · MemValue ( i ) ( 8 ) Ms = LayerNorm ( Hs +WUValue ( m ) ) ( 9 ) where WU is a learnable matrix of shape d× dV .

Summary: The authors propose to extend the standard Transformer architecture by allowing it to attend over factual information, represented as a large memory of dense representations of entity mentions. The resulting architecture is made up of two parts: a mention encoder, which is used to build up the "mention memory"; and the transformer model augmented with attention over the memories. These are pre-trained in two stages for efficiency reasons. The model is evaluated on two claim verification datasets and two QA datasets, showing convincing performance against comparable methods.