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 paper proposes TOME, a Transformer that additionally performs cross-attention over contextual mention encodings (in 1 or 2 layers - TOME-1/TOME-2) along with the usual self-attention. It can be viewed as extending EaE (Fevry et al., 2020) from entities to mentions (aka., virtual KB - Dhingra et al., 2020). TOME is pretrained on 150m mentions extracted from Wikipedia by masked language modeling and other auxiliary objectives. It significantly outperforms EaE on fact checking (HoVer, FEVER) and entity-only QA (TriviaQA, CWQ), but lags behind explicit retrieve-and-read models like RAG, REALM, FiD.