Mention Memory: incorporating textual knowledge into Transformers through entity mention attention
Authors: Michiel de Jong
Yury Zemlyanskiy
Nicholas FitzGerald
Fei Sha
William W. Cohen
Venue: ICLR 2022 Poster (2022)
Keywords: NLP
Entities and Relations
Memory

Abstract: Natural language understanding tasks such as open-domain question answering often require retrieving and assimilating factual information from multiple sources. We propose to address this problem by integrating a semi-parametric representation of a large text corpus into a Transformer model as a source of factual knowledge. 
Specifically, our method represents knowledge with ``mention memory'', a table of dense vector representations of every entity mention in a corpus. The proposed model - TOME - is a Transformer that accesses the information through internal memory layers in which each entity mention in the input passage attends to the mention memory. This approach enables synthesis of and reasoning over many disparate sources of information within a single Transformer model. 
In experiments using a memory of 150 million Wikipedia mentions, TOME achieves strong  performance on several open-domain knowledge-intensive tasks, including the claim verification benchmarks HoVer and FEVER and several entity-based QA benchmarks. We also show that the model learns to attend to informative mentions without any direct supervision.  Finally we demonstrate that the model can generalize to new unseen entities by updating the memory without retraining.

Published as a conference paper at ICLR 2022
MENTION MEMORY:
INCORPORATING TEXTUAL KNOWLEDGE INTO TRANSFORMERS THROUGH EN- TITY MENTION ATTENTION
Michiel de Jong"! Yury Zemlyanskiy" ; Nicholas FitzGerald;; Fei Sha', William Cohen University of Southern California, GGoogle
ABSTRACT
Natural language understanding tasks such open-domain question answer- ing often require retrieving and assimilating factual information from multiple sources. We propose address this problem by integrating a semi-parametric representation of large text corpus into a Transformer model as source of fac- tual knowledge. Specificallyy, our method represents knowledge with "mention memory", table of dense vector representations of every entity mention in a cor- pus. The proposed model TOME is Transformer that accesses the information through internal memory layers in which each entity mention in the input passage attends to the mention memory. This approach enables synthesis of and reasoning over many disparate sources of information within single Transformer model. In experiments using a memory of 150 million Wikipedia mentions, TOME achieves strong performance several open-domail knowledge-intensive tasks, including the claim verification benchmarks HoVer and FEVER and several entity- QA benchmarks. We also show that the model learns attend informative mentions without any direct supervision. Finally we demonstrate that the model can generalize new unseen entities by updating the memory without retraining.
INTRODUCTION
Neural models have greatly advanced the state of the natural language processing and gener- ation tasks. Accordingly, there has been increasing interest in applying neural language models to tasks which require extensive world knowledge to solve (Petroni et al. 202T). Much of this world knowledge can be found distributed over text corpora, which raises the question whether language models pre- pre-trained on text corpora capture this information. Recent work suggests that while lan- guage models may successfully predict facts about the world (Petroni 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 virtual knowledge base (VKB) (Dhingra et al. 2020, Sun et al. 202I). 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 incorporate уKB into a language model by using as an external memory, performing attention over the entire VKB within 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 usefuu for Transformer model, and construct a Mention Memory from 1 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 T).
We argue that the Mention Memory approach has several appealing properties. First, TOME retrieves entity mention representations corresponding to specific entity attributes relations described in the
'Equal contribution. Correspondence to muadejoong.uury.zemlyaniii @usc.edu. Work primarily done Google Research. Code is released https: / / thub. ccm/ googl reaear caaeearchhhaee chaaarggagaaaee///e
tree/master/1ant wwtp://///aanaagg/meeeeeeee tree//maatter//aanaaeeeeenn
onmemory

---

Published as a conference paper at ICLR 2022
TOMEBlock
MentionMenory
TransformerBooc
150M Mikipedia entity mentions releeneeeeen
MenKey
Mervalue
[Perseus] was a great Greek hero
Perseus was great [Greek] hero...
Pre-trained KentionEncdder
PenoryAttentionnlyyy
[Medusa] was slain by Perseus Medusa was slain by [Perseus] (H Simpsonj fictional character
InitialTrans worseroooorrroooo OrwerBloc Block
What the [nationality] the [horo] who killdd [Medusa]?
Figure I: Overview of Mention Memory. pre-trained mention encoder is used to generate dense representations for each entity mention in Wikipedia (approximately 150 million total) which stored a table. The TOME model takes a passage annotated with entity mention boundaries as input, and applies Transformer block. Next, the TOME model applies one or more TOMEBloCKs ocks. Each TOMEBlock contains a memory attention layer and Transformer block.
corpus. This retrieval is much more fine-grained than aggregate entity retrieval methods such as En- tities as Experts (EaE) (Fé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 Second, TOME retrieves dense repre- sentations, which are easy to incorporate into a Transformer model without reprocessing the input, unlike raw text. Therefore, TOME is able 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-reat (Zhao ett al., [202]. 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 semi-parame parametric, so knowledge can be added or updated by applying the mention encoder to new text without retraining.
In order verify the model's 's capacity to capture accurate factual information the corpus, we start by evaluating TOME the Ho Ver (Jiang et al. 2020). FEVER (Thorne et al. 2018) 8) and FM2 (Eisen- schlos al. 202T) claim verification datasets, on which strongly improves performance over en- tity aggregate and comparable retrieve-and -nddrdd baselines. We demonstrate that the model learns to attend to informative mentions for verifying claims using only the verification accuracy as a sig- nal. Ablations show the memory is crucial for performance, and that the model effectively use larger memory than was pre-trained on. In a second set of experiments we evaluate TOME on question-answering benchmarks TriviaQA (Joshi al.) (017)), Complex CempleexeeQuusttonn (Talmor & Berant 2018) 8) and EntityQuestions (Sciavolino et al. (202)), improving performance over compara- ble baselines. Finally show that the model can be adapted to generalize 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 the corpus. Every time an entity appears in passage - "/Barack Obama) was elected president in 2008" some property of the entity or its rela- tion other entities described. The first component of our method, the Mention Encoder model, responsible for distilling information from entity mentions in the corpus into high-dimensional men- tion 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 Transformer. The second component of our method, the TOME model, applies sparse attention over the Mention Mem- ory to incorporate external information from the corpus into a Transformer model. An overview of the whole method is shown in Figure
2

---

Published as a conference paper at ICLR 2022
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, propose train the models in two stages. First, 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 En- coder 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 the model is a passage x Z1,.... , PT 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 [Esttrt]] and [Fena] 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 the [Estrtt] nationality [Lena]] of the [Estort] hero [Eena] who killed [Estar]] Medusa [Enad]" Each mention m in a passage is described by a tuple (s, e), where s and e are start and end positions the mention. We consider entity markers to be part of the corresponding mention, so that x, [Estard] and De = [Eena]- Representations of these tokens are later used to generate mention encodings.
2.1 CONSTRUCTING MENTION MEMORY FROM CORPUS
2.1.1 MENTION ENCODER
Let И E RT'x be token representations where d is the hidden dimension, such that H, e Rd is the contextualized embedding for the i-th token. Following Févry et al. (2020) we compute the encoding of span (s, e) as learnable linear projection W of the concatenation of its start and end token representations H, and He
(1) 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 SpanEncodingLayerr do not share weights.
SpanEncodingLayerrs H, (s) W HH; H.]
The Mention Encoder is a Transformer model with two final SpanEn
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 linked mentions, which we hypothesize contain information that can be retrieved and grounded. We gather mention encodings into matrices MemKey E R xdx and MemValue E Nxdy where N is the total number of linked entity mentions in English Wikipedia (approximately 150 million) and dæ and dy are dimensions of key and value encodings. Additionally, we record entity (Wikipedia) IDs of mentions in MemEnt R which we use as labels for auxiliary losses, not as inputs to the model or supervision on retrieval. Me)Key(i), MemValue 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 text corpus into Transformer by applying sparse attention over the Mention Memory. The model consists of one or more TOKEBlocks, each contain- ing 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 ac- cess 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(F]
H' TransformerBlockx.1
(3)
3

---

Published as a conference paper at ICLR 2022
The TOME model encodes input passages with the word embedding layer and initial Transformer block and then applies one or more TOMEBlocks
H0 Inttaaltraattorrrfiioooooooooaar
(4)
H TOMEBLock((1 1...
(5)
In this work consider two configurations of the TOME model: TOME-1 and TOME-2, with one and two JOMEBlocks respectively. Each TOMEBlock of TOME-2 contains half as many Transformer layers in TOME-1 hold the total number of Transformer layers fixed between models.
2.2.1 MemoryAttention
Each memory attention layer is implemented sparse dot-product attention layer that takes the output H of the previous Transformer block, incorporates information from the Mention Memory, and returns representation M (omitting layer indices). Consider mention that starts at position s and ends at position& e. We start by computing its query mention encoding Query(m) by applying SpanEncodingLayer
Query(m) SpanEhcodingLayerr(H,, (s,))),
(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, 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 (aurry(m),, We compute attention over these memories and incorporate the result into the token contextual representation at position s
0g exp(Query(m) MemKey(i)), i E TopMem((uery(m))
Value(m) =
Σ
0r MemValue
(8)
iE TepMonn(Quory(m))
M, LayerNorm Wy WyValue(m))
where Wy is learnable matrix of shape d x dv.
2.2.2 SPARSE LARGE-SCALE RETRIEVAL.
Approximate nearest neighbor search (ANNS) can be performed cheaply using one of multiple ANNS libraries, for example ScaNN (uoo et al. 2020). We implemented two on-device search methods to avoid the engineering complexity of real-time communication with an ANNS server, though we have verified this is also viable. The first naively computes a simple dot-product between passage queries and memory keys, and was used in our main experiments it was easiest to im- plement. We also implemented and will be releasing much faster version based on CPU ANNS methods. The memory is sharded over devices, that the device-memory overhead is negligible.
Holding the number of entries in memory fixed, the compute cost of retrieval from memory does not grow with the size of the reader the dimensionality of the memory values, that the relative cost of the memory layer becomes smaller with reader size. In particular, the overhead from the memory used in our pre-training setting is small for BERT-Large and up. More details on ANNS implementation and overhead can be found in Appendix C
2.3 MENTION ENCODER PRE-TRAINING
While backpropagating through Wikipedia-scale mention memory is challenging, it is possi- ble to train smaller-scale memory architectures end-to-end. We take an approach inspired by MARGE (Lewis et al.) 2020a) and READTWICE (Zemlyanskiy et al. (021)) which apply cross- attention over documents within batch. In particular, we process passages in each batch twice. As a first step, the Mention Encoder model generates mention encodings from each passage and aggregates the mention encodings into a batch-wide memory table. In the second step. apply TOME architecture that attends to the batch memory, which we call BATCH-TOME. Note that BATCH- TOME is just used for pre-training the Mention Encoder and not evaluated on any downstream tasks.
4

---

Published as a conference paper at ICLR 2022
Mention Encoder and BATCH-TOME are jointly trained end-to-end that the Mention Encoder is encouraged to produce mention encodings that contain useful information for BAMCH-TOME.
We want to make the batch memory contains relevant mentions, we pre-train the models on batches of passages constructed from related Wikipedia articles with high entity overlap. Ap- pendix A. provides more details Mention Encoder data generation. We the pre-trained Mention Encoder to construct the Mention Memory table from corpus, and use the BATCH-TOME model the initialization point for TOME-specific pre-training (described in Section 2.4..
Masked language model. Our primary pre-training objective is the standard masked language modeling task, with the loss computed based on the output of the second read (BATCH-TOME). To encourage the model to rely on memory, we increase the task's 's difficulty relative to standard BERT pre-training by masking entity mention tokens more aggressively.
Coreference resolution. We wish to encourage the Mention Encoder to represent the entity at- tributes expressed by entity mentions, so we also employ an -riented pre-training task the output of BATCH-TOME for which such attribute information is likely to be especially helpful. Un- like Entities as Experts (Févry et al., 2020), BATCH-TOME does not use entity embeddings, so we cannot use the entity linking task. Instead, we apply related entity coreference resolution objective, which asks the model to predict whether two linked mentions correspond to the same entity based on the similarity of their encodings. Given that entity surface forms are frequently masked, the model needs to instead use the properties of other mentions in the batch to determine which entity it is most compatible with, incentivizing the Mention Encoder to encode such properties. We com- pute a coreference mention encoding for every linked mention in the batch by applying a separate SpanEncod ing.anynednnlyyrrrrr the output of BATCII-TOME. The loss is implemented using cross-entropy over dot-product similarity scores. See Appendix A.2 for details.
2.4
TOME PRE-TRAINING
As TOME attends to the full Mention Memory instead of in-batch memory, we do not employ the batching procedure from Mention Encoder pre-training, instead sampling Wikipedia passages ran- domly. For the same reason, we replace the in-batch entity coreference objective by Mention Mem- ory entity coreference, in which the model has to predict which mentions from the Mention Memory share an entity with the input mention. The goal of this auxiliary objective is to incentivize the model to learn retrieve informative mention encodings to solve the semantically challenging task. Men- tion Memory entity coreference also allows us to solve tasks like TriviaQA or Complex without a decoder by directly predicting the answer entity.
Entity prediction. Analogous to batch coreference resolution loss we compute mention encoding 2m using the output of the TOME model. As in section 2.2, TopMem( (zm) returns the top K memories with the largest dot product between the mention encodings 2m and key mention encodings MemKey from the Mention Memory. The score EntProb(m, j) of entity j equals the sum of attention weights of memories corresponding to this entity.
EntProb(m..)) EieTopMen(t..)) exp(2m 2m MemKey (i) (MemEnt(i) =
(10)
Σ [TTopteea(=m)) (xm MemKey(i))
The final entity prediction is arg max, EntProb(m,j). Entity prediction loss (m)(m) for mention m of entity Ent(m) is Lep (m) log EntProb(m,, Ent(m)). Total loss equals the average loss over linked input mentions for which at least one memory of the same entity is retrieved.
Disallowed same passage retrieval. For each passage in the pre-training corpus, there exist mem- ories corresponding to mentions in the passage generated from the unmasked version the same passage. In order to prevent the model from "cheating' by attending to such memories, we set the attention weight for all memories from the same passage to zero.
3 RELATED WORK
Our approach lies at the intersection of three lines of work: i) knowledge-uugmented language models, ii) employing text corpus as a virtual knowledge base, iii) andried methods.
5

---

Published as a conference paper at ICLR 2022
Knowledge-augmentee language models. Entities Experts (EaE) (Fvry et al. 2020) injects in- formation into Transformer model model with an intermediate attention layer over trainable entity embeddings, which serve as aggregate representation of entity information in text corpus. In contrast, TOME attends to a much larger table of mention encodings, allowing for retrieval of more fine-grained information. Attending to mentions as opposed entity representations also enables TOME to generalize to unseen entities. FiLM (Verga et al. (021] extends EaE by adding an attention layer over facts from a KB on the output of the Transformer. The fact attention layer enables more fine-grained queries but still retrieves aggregate entity embeddings as values, which are also not reasoned over by the Transformer. KnowBERT (Peters al.) 2019 is similar to EaE, but with entity embeddings generated from a KB instead of trained end-to-end with text corpus. MARGE (Lewis) jet al., 2020a) and READTWICE (Zemlyanskiy et al. 202T) incorporate dense representations from other passages within the same batch into a Transformer through sparse top-k attention. The first pre-training stage of our method for training the Mention Encoder is similar to MARGE and READ- TWICE. However, TOME performs global attention over full corpus, rather than a single batch. Furthermore, TOME attends to a Mention Memory consisting of pre-computed dense representa- tions. Therefore TOME is not limited to downstream task with batches of relevant documents, and does not need to apply an expensive reader model to an entire batch of documents for each input.
Text corpus virtual knowledge base. DrKIT (Dhingra et al. 2020) performs multi-hop question answering by using a text corpus as a virtual knowledge base. Similar to TOME, the authors apply a mention encoder to convert the corpus into a table of mention encodings. A Transformer model encodes the question into dense queries, which are compared with the mention encodings to fraverse the VKB. Conversely, TOME retrieves mention encodings, and then jointly processes them inside the Transformer. In follow-up work to DrKIT, OPQL (Sun et al. 2021) uses a FiLM-like approach to access memory of relation mentions, which are encoded with a sule-vupe relation encoder. However, the relation mention encoding combines mention-specific relation representation with EaE-like entity encodings, so they are less fine-grained than TOME'S encodings. Unlike TOME, OPQL also lacks sparse large-scale retrieval mechanism. and relies on ad hoc heuristics to limit the size of the memory [ MOLEMAN (FitzGerald et al. 2021) compares a passage mention encoding with mention encodings from a corpus to perform entity linking, but does not retrieve the mentions.
Retrieve-and-reat methods. REALM (Guu et al. 2020) learns to retrieve relevant passages from text corpus in sulf-supervised manner. Retrieved passages concatenated to the input passage which is then re-encoded by a Transformer model to perform a task. The key difference between retrieve-and-read approaches (Guu et al. 2020; Karpukhin al. 2020 Lewis et al. 2020b Izac- ard & Grave 202T) and TOME is that TOME retrieves dense representations, as opposed to text. That means that TOME only applies a reader model once single input, while retrieve-and-read approaches have to apply an expensive BERT reader to many different passages. In addition, Trans- former models can only process relatively short sequences, which imposes a binding constraint on the number of retrieved text passages that can be processed together, whereas TOME can retrieve and reason over information from many sources inside the same reader. Generative models like RAG (Lewis et al. (220bb) or FiD (Izacard & Grave, (021) attend to different retrieved documents in the decoder, but still have to apply a BERT read for every retrieved document, do not consider interaction between retrievals while encoding the question, and cannot perform iterative retrieval.
4 EXPERIMENTS
4.1
EXPERIMENTAL SETUP
The Mention Encoder is based on BERT-base model with two final SpanFacodingLayers that produce key and value encodings. Mention Encoder and BATCH-TOME share Transformer weights during Mention Encoder pre-training. The Mention Memory consists of mention encodings for N = 150 million linked Wikipedia entity mentions. Transformer layers in TOME and BATCH-TOME models are equivalent to those in the BERT-base model. The TOME contains 4 Transformer layers. TOME-1 has single TOMEBlock with 8 Transformer layers, and TOME-2 has two TOMEBloCks ocks with 4 Transformer layers each. Therefore, the number of trainable parameters in 1 1 and TOME-2 2 is approximately the same as in BERT-base. We use smaller
It should be noted that absent heuristics, the number of potential relation mentions (i.e., entity mention pairs) is much larger than the number of entity mentions.

---

Published as a conference paper at ICLR 2022
Table 1: Accuracy on claim verification datasets. #Encoded refers to the number of passages en- coded by BERT reader to answer a single question.
Model
#Params
#Encoded
Horers
FEVERtest
FM2dev
RAG
620M
100
72.5
REALM
330M
5
66.1
67.1
65.8
Entities as Experts
360M
1
66.6
63.6
63.5
TOME-1
220M
1
72.8
67.8
67.7
TOME-2
220M
1
73.1
68.1
68.4
Mention Memory containing 38m uniformly sampled memories for TOME pre-training. During fine-tuning and evaluation utilize the full Mention Memory. Appendix A contains more details.
4.2 BASELINES
We compare TOME with existing methods that utilize textual information from corpus in lan- guage model. These can be divided into generative LLMs (T5), entity embedding retrieval (Entities Experts, OPQL), extractive redieeveaandrrad (REALM) and generative retrieve-and and-red (RAG, Fusion-in-Decoder)) TOME occupies a novel position in the space of retrieval models, being more fine-grained than entity embedding retrieval methods, but performing all its reasoning with single BERT read, unlike eetriee-aand-res methods. The most closely comparable models are Entities as Experts and REALM, and we use these as our primary baselines. We report other baselines for reference, with the caveat that these results are not apples-to-apples: RAG and Jusion-in-Decoder have large decoders and retrievers and encode a large number of passages with a BERT reader for each question compared to TOME's single read. Fusion-in-Decce supervision for retrieval. We mark the number parameters and BERT applications for each base- line in the result tables. Consistent with retrieve-and-read, we count the parameters of the Mention Encoder and TOME, but not the size of the non non-trainable and sparsely accessed Mention Memory.
and RAC also use ground-truth
4.3 CLAIM VERIFICATION
Data. Our first set of experiments evaluates TOME on the claim verification tasks FEVER (Thoree et al.) 2018), Ho (Jiang et al., 2020), and FM2 (Eisenschlos et al., 202I in which the model is provided with a claim and has to determine whether the claim is supported by the Wikipedia corpus. FEVER is larger dataset with 186k claims for which most of the claims be verified with a single Wikipedia passage. In contrast, HoVer is smaller with 26k claims, but is explicitly constructed require evidence from multiple sources and multiple reasoning steps. FM2 is also smaller and is constructed through an adversarial game that leads to more challenging retrieval. The claim verification training data contains gold evidence passages, but unlike most published results we do not use these, leaving only the accuracy of the claim verification to guide the retrieval.
Results. Table D contains our claim verification results. TOME outperforms both Entities as Ex- perts and REALM, especially on HoVer and FM2. This is consistent with the properties of TOME: HoVer requires combining detailed information from multiple sources, which TOME especially well equipped to do compared to aggregate entee or aedrieveaand-red models. FM2 features generally challenging retrieval and may benefit from contextualizing retrieved evidence.
4.4 QUESTION ANSWERING
Data. In a second set experiments we evaluate TOME on TriviaQA (TQA) (Joshi et al. 2017).
ComplexWebee WebQuestions (CWQ) (Talmor & Berant, 2018) and EntityQuestions (EQ) (Sciavolino ct al. 202T), open-domain QA tasks for which most answers are Wikipedia entities. We approach these datasets entity-linking tasks, Févry et al. (2020). We append mask token to each ques- tion, which is marked as a question mention. The probability for each candidate entity is predicted the aggregate attention weight on mentions of the entity (Section 2.4p. Questions with answers that do not correspond to entities in our entity vocabulary are marked as answered incorrectly. TQA
ARAG is initialized from DPR which is trained with gold retrieval passages for TriviaQA.

---

Published as a conference paper at ICLR 2022
Table Accuracy on open-domain QA datasets TriviaQA (TQA), ComplexWeb (CWQ) and
(EQ). #Encoded refers to the number of passages encoded by BERT reader to answer a question. devee corresponds to TQA with train and dev samples limited to those with Wikipedia entity as answer. See Appendix B.3 for full results.
Model
#Params
#Encoded
TQA
TQA
TQA
CWQdev
EQdev
RAG
620M
100
56.8
68.0
Tusion-in-Decodee
440M
100
65.0
77.1
REALM
330M
5
55.8
67.1
63.4
46.7
59.0
T5-3B
3B
38.7
T5-11B
11B
42.3
50.1
Entities as Experts
360M
43.2
53.4
51.3
42.7
32.5
OPQL
220M
41.1
TOME-1
220M
50.8
61.1
60.3
44.9
62.1
TOME-2
220M
54.6
65.8
64.8
47.7
66.0
Table TOME-2 retrievals for the second HoVer dev sample. We show top-1 retrieval results for the first ( (-+1) memory attention layer for two passage mentions. Memory mentions are in brackets.
Claim: Greater Swiss Mountain Dog and Harrier are both dog breeds. Label: TRUE
Greater Swiss Mountain Dog > Breed History the origin of the [Greater Swiss Mountain Dog] is not definitively known.
Harrier -1 The harrier is a medium-sized dog breed of the [hound] class, used for hunting...
consists of 96k trivia questions, for which 84% of answers correspond to Wikipedia entity. We use the open-domain setting without gold evidence passages. In order to compare head-to-head per- formance, we also report results subset of TQA with only questions with Wikipedia entities as an answer. CWQ consists of 35k complex questions (compositions, conjunctions, etc.) for which 94% of answers correspond to Wikipedia entity. EQ contains challenging questions involving rare entities, with Wikipedia entities as answers.
Results. Table 2 contains the results for TQA. CWQ and EQ experiments. Like TOME, Entities as Experts and OPQL treat the above datasets as entity-linking tasks. REALM performs extractive QA, while T5, RAG and Fusion-in-Decode generate the answer. We note similar pattern of results as for claim verification. TOME strongly outperforms Entities as Experts on all tasks. TOME performs slightly better than REALM a simple task like TriviaQA (entity subset) and strongly outperforms REALM on more challenging tasks that require multiple (CWQ) or challenging (EQ) retrieval.
4.5 QUALITATIVE PROPERTIES OF TOME
What memories does TOME retrieve? Given that TOME retrieval is unsupervised, it is natural to ask what memories it leams to retrieve. First, we observe that BATCH-TOME and TOME trained on just the MLM objective learn attend to memories of the same entity the passage linked mention (55% and 41% % average attention score). This is promising entity mentions from the same entity often contain mutually relevant information. Quantitative evaluation of downstream retrieval is challenging TOME often retrieves mentions that are not part of, but equally informative gold passages. Instead, we provide TOME retrievals on the first three samples of the HoVer dev set to demonstrate its retrieval behavior without cherry-pickin Table (aeeeddemonn demonstrates successful simple retrieval, while Table 4 displays interesting multi-hop retrieval. The last found in Appendix D
Importance of memory size. Figure shows claim verification performance as function of memory-size during fine-tuning (pre-training memory size held constant). For smaller mem- ory sizes, entries in memory uniformly sampled from the full Mention Memory. Performance increases smoothly with memory size. Larger memory size yields diminishing returns, perhaps reflecting that entity mentions may contain overlapping information.
8

---

Published as a conference paper at ICLR 2022
Table 4: TOME-2 retrievals for the first HoVer dev sample. We show top-1 retrieval results for the first (- and the second (->2) memory attention layers for passage mentions "Life Goes On" and "Hungry" Memory mentions are in brackets. The first retrieval for the "Life Goes On" is a different song with the same name and the first retrieval for "Hungry" is related but not useful. However, the second retrieval for "Life Goes On" identifies the correct song and describes its position on the album while the second retrieval for "Hungry" captures its position relative to "Life Goes On".
Claim: The song recorded by Fergie that was produced by Polow Da Don and was followed by Life Goes On was Hungry. Label: TRUE
Life Goes On - and Johnny produced the chart topping hits "All Bout U", "How Do Want It" and L'fe Goes On"l....
Life Goes On -2 On November 11, 2016, Fergie released the third single from the album, L'ifee Goes On")...
Hungry -1 Polow da Don, is an American record producer, songwriter and rapper. His cousin is [Atlanta] singer Monica. Jones has produced variety of singles for multitude of artists including "Anaconda" by Nicki Minaj (2014), "Love In This Club" by Usher (2008), "Buttons" by the Pussycat Dolls (2006), "Hungry" by Fergie
Hungry - "Life Goes On" is a song recorded by American singer Fergie for her second studio album, Double Dutchess (2017). ... The song serves as the third single from [Fergie's] second studio album, following "Hungry""
74
Table 5: Accuracy on held-out subset of Trivi- aQA and ComplexWebQuestion tions. TOME- WebQuestions (CWQ) ques- was pre-trained and fine- tuned with memory without entities from held- out set and evaluated with full memory. Note that performance is considerably lower than on the full dev set as answers in the held-out set (which are in dev but not train) are more likely to be rare entities.
Accuricy %
68
Ho (dev) FEVER (dev)
64
50
100
150
Dataset
TriviaQAde:
CWQqte
Memory size during fine-tuning, millions.
TOME-1
17.4
16.4
Figure 2: Claim verification accuracy func- tion of fine-tuning memory size (in millions).
TOME--1unsees
17.6
16.7
Zero-shot transfer to unseen entities. An important advantage of memory architectures is that the behavior of the model can be steered by deciding what include the memory. Here we show that the TOME model can use information that was not present in memory during training. We sample questions the TQA and CQA dev sets, and generate a subset of the memory without any mentions corresponding to the answer entities for those questions. Then we pre-train and fine-tune model on this smaller memory, which we call TOME-unseen. We evaluate TOME-unseen on the sampled questions using the full memory for evaluation only, and compare to standard TOME. Table 5 shows that using full memory only during evaluation does not lower performance.
5 CONCLUSION
We introduced TOME, a Transformer model that performs attention over a semi parametric repre- sentation of the entire Wikipedia text corpus. This representation, or Mention Memory, consists of dense encoding for each entity mention in Wikipedia. TOME can retrieve information from multi- ple sources without supervision, aggregate information within the Transformer, and reason over the retrieved information. TOME leads to strong improvements multiple open- ppen-domain claim verifica- tion and entity-based question answering tasks.
We replaced the original song title with the song "Hungry" as the original may he inappropriate.

---

Published as a conference paper at ICLR 2022
Ethics. We pre-train and evaluate public data without privacy information. We aim to improve performance on established academic NLU tasks, which may eventually extend to per- formance tasks with more direct ethical implications. Improving the ability of Transformers to incorporate textual knowledge may help with reducing factual errors and hallucinations for large language models, which known to carry risks.
Reproducibility. We have release coded for the model and pre-training, as well as pre-trained TOME-1 and TOME-2 model checkpoints and corresponding Mention Memory. Appendix A pro- vides details of the pre-training procedure. The most challenging part of the model to reproduce is the sparse attention layer, described in Appendix We do not release pre-training data, but the data can be easily reproduced from Wikipedia dumps and the publicly available Google Cloud API.
REFERENCES
Livio Baldini Soares, Nicholas FitzGerald, Jeffrey Ling, and Tom Kwiatkowski. Matching the blanks: Distributional similarity for relation learning. In ACL 2019, 2019.
Boxi Cao, Hongyu Lin, Xianpei Han, Le Sun, Lingyong Yan, Meng Liao, Tong Xue, and Jin Xu. Knowledgeable or educated guess? revisiting language models knowledge bases. In ACL/JJCNLP 2021, 2021.
Bhuwan Dhingra, Manzil Zaheer, Vidhisha Balachandran, Graham Neubig, Ruslan Salakhutdinov, and William W. Cohen. Differentiable reasoning over virtual knowledge base. In ICLR 2020, 2020.
Julian Martin Eisenschlos, Bhuwan Dhingra, Jannis Bulian, Benjamin Börschinger, and Jor- dan Boyd-Graber. Fool me twice: Entailment from wikipedia gamification. arXiv preprint arXiv:2104.04725, 2021.
Thibault Févry, Livio Baldini Soares, Nicholas FitzGerald, Eunsol Choi, and Tom Kwiatkowski. Entities experts: Sparse memory access with entity supervision. In EMNLP 2020, 2020.
Nicholas FitzGerald, Daniel M. Bikel, Jan A. Botha, Daniel Gillick, Tom Kwiatkowski, and Andrew McCallum. MOLEMAN: mention-only linking of entities with mention annotation network. In ACL-IICNLP 2021,
Ruiqi Guo, Philip Sun, Erik Lindgren, Quan Geng, David Simcha, Felix Chern, and Sanjiv Kumar. Accelcrating large-scale inference with anisotropic vector quantization. In ICML 2020, 2020.
Kelvin Guu, Kenton Lee, Zora Tung. Panupong Pasupat, and Ming-Wei Chang. Retrieval augmented language model pre-training. In ICML 2020, 2020.
Gautier Izacard and Edouard Grave. Leveraging passage retrieval with generative models for open domain question answering. In EACL 2021, 2021.
Yichen Jiang, Shikha Bordia, Zheng Zhong, Charles Dognin, Maneesh Kumar Singh, and Mohit Bansal. Hover: A dataset for many-hop fact extraction and claim verification. EMNLP 2020, 2020.
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with gpus. arXiv preprint 1rXiv:1702.0873.1 2017.
Mandar Joshi, Eunsol Choi, Daniel S. Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In ACL 2017, 2017.
Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick S. H. Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. In EMNLP 2020, 2020.
Mike Lewis, Marjan Ghazvininejad, Gargi Ghosh, Armen Aghajanyan, Sida I. Wang, and Luke Zettlemoyer. Pre-training via paraphrasing. In NeurlPS 2020, 2020a.
10

---

Published as a conference paper at ICLR 2022
Patrick S. H. Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, Sebastian Riedel, and Douwe Kicla. Retrieval-angmmened generation for knowledge-intensive NLP tasks. In NeurlPS 2020, 2020b.
yya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In ICLR 2019, 2019.
Matthew E. Peters, Mark Neumann, Robert L. Logan IV, Roy Schwartz, Vidur Joshi, Sameer Singh, and Noah Smith. Knowledge enhanced contextual word representations. In EMNLP-IJCNLP 2019.
Fabio Petroni, Tim Rocktäschel, Sebastian Riedel, Patrick H. Lewis, Anton Bakhtin, Yuxiang Wu, and Alexander H. Miller. Language models as knowledge bases? In EMNLP-IJCNLP 2019, 2019.
Fabio Petroni, Aleksandra Piktus, Angela Fan, Patrick S. H. Lewis, Majid Yazdani, Nicola De Cao, James Thorne, Yacine Jernite, Vladimir Karpukhin, Jean Maillard, Vassilis Plachouras, Tim Rocktäschel, and Sebastian Riedel. KILT: benchmark for knowledge intensive language tasks. In NAACL-HLT 2021, 2021.
Christopher Sciavolino, Zexuan Zhong, Jinhyuk Lee, and Danqi Chen. Simple entity-centric ques- tions challenge dense retrievers. In Marie-Francine Moens, Xuanjing Huang, Lucia Specia, and Scott Wen-tau Yih (eds.), EMNLP 2021, 2021.
Haitian Sun, Patrick Verga, Bhuwan Dhingra, Ruslan Salakhutdinov, and William W. Cohen. Rea- soning over virtual knowledge bases with open predicate relations. In ICML. 2021, 2021.
Alon Talmor and Jonathan Berant. The web as a knowledge-base for answering complex questions. In NAACL-IILT 2018, 2018
James Thome, Andreas Vlachos, Christos Christodoulopoulos, and Arpit Mittal. FEVER: large- scale dataset for fact extraction and verification. In NAACL-HLT 2018, 2018.
Pat Verga, Haitian Sun, Livio Baldini Soares, and William Cohen. Adaptable and interpretable neural memory over symbolic knowledge. In ACL 2021, 2021.
Yury Zemlyanskiy, Joshua Ainslie, Michiel de Jong. Philip Pham, Ilya Eckstein, and Fei Sha. Read- twice: Reading very large documents with memories. In NAACL-HLT 1021,2021
Chen Zhao, Chenyan Xiong, Jordan L. Boyd-Graber, and Hal Daumé III. Multi-step reasoning over unstructured text with beam dense retrieval. In NAACL-HLT 202/, 2021.
A PRE-TRAINING
We train English Wikipedia, processed with the entity linking and named entity recognition tools from the Google Cloud NLP AP We use existing hyperlinks in Wikipedia as additional entity annotations. All models are pre-trained on 128 TPUs using AdamW optimizer (Loshchilov & Hutter 2019) with learning rate lc-4 and batch size of 4096. Each passage in the batch has length 128, excluding entity tokens. The Mention Encoder and BATCH-TOME are pre-trained for 1 million steps with 50k warmup steps, and TOME is trained for 500k additional steps with 25k warmup steps after initialization from BATCIITOMME Both models are trained with linear learning rate decay. Mention Encoder and BATCH-TOME share Transformer weights during Mention Encoder pre-training. We apply gradient clipping with a norm of 1.0 and weight decay 0.01. Weight decay applied to all weights except layer norm and bias weights.
BATCH-TOME TOME and TOME are trained with weight 0.85 on the MLM objective and 0.15 on the entity coreference resolution objective. We mask 20% of whole entity mentions and 10% of other tokens. We limit the coreference resolution objective to mentions of the million most frequent Wikipedia entities. We use 24 mentions per sample, with batch size of 32 samples per TPU. We subsample mentions uniformly if the average number of annotated mentions on a TPU exceeds 24. Key mention encoding have dimension dx 128 and value and coreference mention encodings have dimension dy dc = 512.
https://cll oud. google www.ccom/nauut hag.auurralllaaguuuggggooee//aae
ty_analyais

---

Published as a conference paper at ICLR 2022
Disallowed same passage retrieval for Mention Encoder. We want the model to memory as a source of additional information for processing a passage. Therefore, we explicitly set attention weights to 0 for memories generated from the same passage the current one.
A.1 MENTION ENCODER DATA GENERATION
We pre-train Mention Encoder to produce mention encodings that are useful for BATCH-TONME. In order to provide BATCH-TOME with an incentive to the memory, need to ensure that mentions from different samples within a batch are relevant to each other. We achieve this by batching passages from the or related Wikipedia articles.
We generate clusters of 256 passages from Wikipedia articles using greedy method. First, we create cluster from the longest unused Wikipedia article and add related articles until the cluster consists of 256 passages. In particular, at each step we add the article with the largest Jaccard similarity between its entity set and the entity set of articles in the current cluster.
A.2 COREFERENCE RESOLUTION LOSS
For every linked mention m in the batch we compute a mention encoding 2m by applying separate SpanEncod ingLayer on the output of the BATCH-TOME, First, we compute the loss for every linked mention m in the batch. To this end, we denote linked mentions in every other passage in the batch as positive, P (m), if they have the same entity ID and negative, P (m), otherwise. The loss per mention is an average of cross-entrope losses per every positive mention m P (m)
exp(2 2m
Leoref(m)
log Σ
exp( - =m+)
Em EP (m) exp 2 =m-)
37 EMTEP+(m) (m)
The total loss is average of losses per linked mentions which have at least one positive mention (set P+ (m) is not empty).
B EXPERIMENTS
B. FINE-TUNING SETUP
TOME is fine-tuned 32 TPUs using the Adam optimizer with learning rate of le-5 and total batch size 32. In contrast to pre-training we set max mentions to 32 per sample for fine-tuning. We 1000 warmup steps and linear learning same as during pre-training. We take the highest scoring checkpoint on dev sets and evaluate on the test We use the spaCy noun chunker detect noun phrases and treat these as claim/question entity mentions.
decay. Gradient clipping and weight decay the
The model can be fine-tuned with full memory on server of 8 A100 GPUs or 16 v3/v4 TPUs. A model with half memory (75M mentions) can be fine-tuned on 8 V100s/P100s or 8 TPUs.
B.2 BASELINES
Following Guu al. 2020 we used REALM perform extractive question answering on TriviaQA and Complex WebQuestion datasets. We also adapted the model to the classification setting in order apply to claim verification tasks. Given an input claim X we compute probability of a prediction Y (whether the claim holds true or not) as a marginalization over retrieval Z.
Pr PY Z)
Σ Pr(Y 2) Pr x
zeZ
where Pr(Y|X, Z = 2) is the output probability produced by the reader model and Pr = 2 produced by the retrieval model.
B.3 CLAIM VERIFICATION
See Table for the results on development and test splits claim verification datasets. Additionally, Table comppares our FM2 results the original dataset baselines.
12

---

Published as a conference paper at ICLR 2022
Table 6: Accuracy on claim verification datasets. #Encoded refers to the number of passages en- coded by BERT reader to answer a single question. EaE stands for Entities as Experts model.
Model
#Params
#Encoded
Ho Vereer
Ho Ve
FEVER
FEVER
FM2d
RAG
620M
100
74.5
72.5
REALM
330M
5
67.3
66.1
70,4
67.1
65.8
EaE
360M
66.2
66.6
66.1
63.6
63.5
TOME-1
220M
73.6
72.8
70.5
67.8
67.7
TOME-2
220M
74.1
73.1
71.1
68.1
68.4
Table 7: Accuracy FM2 compared with original dataset baselines. Oracle refers to oracle retrieval followed by BERT-Bas reader.
Model
Accuracy
Oracle (Fisenschlos et al. 2021)
69.3
DPR (Eisenschlos al. 2021
64.2
EaE
63.5
REALM
65.8
TOME-1
67.7
TOME-2
68.4
B.4 QUESTION ANSWERING
dataset from Sciavolino et al. (02IT). The dataset consists of questions involving rare entities, making it especially challenging for modern retrievals methods such DPR. Evaluation results for TOME models and baselines are shown in Table & and Table 9 Following (Sciavolino et al. (021)) we report recall at 20 as an evaluation metric. Since TOME retrieves mentions rather than passages, direct comparison difficult. We evaluate TOME conservatively, treating recall at 20 as successful if one of the 20 highest scoring mentions belongs the correct entity (in contrast DPR. for which the correct answer only has be somewhere in the retrieved 100 word document).
We report additional results the
TOME sets the state of the art on this dataset and outperforms DPR by very large margin. REALM cannot be fairly compared to DPR due to longer retrieved passages (100 vs 288 tokens). There- fore, we perform a separate experiment using accuracy with REALM as a baseline, showing large performance gains over REALM as well.
B.5 IMPORTANCE OF PRE-TRAINING ORJECTIVES
We perform several ablation experiments for the pre-training procedure (see Table 10). First, results show that the entity prediction objective (c.f. Section 2.4 is not essential for TOME pre-training. Performance on claim verification datasets (FEVER and HoVer) is not affected by whether we use entity prediction for pre-training. More surprisingly, removing this objective only slightly decreases the performance on entity question answering datasets (TriviaQA and Complex WebQuestions). We predict entities for question-answering in the way do for the entity prediction objective during pre-training (c.f. Equation 10 important.
so we expected the entity prediction auxiliary loss to be
Table 8: EntityQuestions recall @0
Table Entity top-1 accuracy
Model
Recall @:20
Model
Accuracy
DPR Sciavolino et 2021
65.4
Entities as Experts
32.5
BM25 Sciavolino et al 2021
71.2
REALM
59.0
TOME-
83.3
TOME-1 1
62.1
TOME-2
83.8
TOME-2
66.0
13

---

Published as a conference paper at ICLR 2022
Table 10: Performance ablations for pre-training objectives experiments.
Dataset
Ho Veraev
TriviaQAde dEVER
CWQqdev
TOME-1
73.6
70.5
50.8
44.9
w/o entity coreference loss
69.8
68.4
42.5
40.5
w/o entity prediction loss
73.7
70.7
49.4
43.8
On the other hand, a related entity coreference objective (c.f. Section A.1 and Appendix (A.2) is
crucial for Batch-TOME and Mention Encoder pre-training. That is consistent with our intuition that semantically challenging tasks incentivize the model to store useful information in memory.
B.6 TOME INITIALIZATION
We initialize TOME model with pre-trained AATCH-TOME model which we find to be especially im- portant for warming up retrieval. If TOME initialized from scratch (or even from BERT weights), TOME does not learn to use the memory. In fact, TOME has to be initialized from the same BATCH- TOME used to generate the memory. This implies that multi-stage training is vital ingredient for TOME to succeed. Our explanation for why TOME is sensitive initialization is that TOME needs to learn two skills: first, to effectively use retrieved mentions for its predictions, and second, to retrieve relevant mentions. Learning both capabilities end to end gives rise mutual dependence: to get a signal for learning how to use retrieved mentions, the retrieved mentions have to be useful, and to learn to retrieve useful mentions, the model needs utilize retrieved mentions. If initialized from scratch, the model is not able to learn both skills simultaneously. The pre-training stage with the smaller in-batch memory functions as curriculum address that problem.
C NEAREST NEIGHBOR SEARCH
Nearest neighbor search is extremely common problem, and there exist numerous approaches and packages for fast approximate nearest neighbor search (ANNS). (Guo et al. 2020 Johnson et al., 2017]. Most approaches employ two methods for fast search: 1) compress the search table through projecting to a lower dimension and quantization and perform comparisons in this compressed space and 2) divide the search table in buckets of similar items, and search only subset of the buckets. Retrieve-and-read use ANNS packages to search for related passages (Guu et al. 2020 Lewis et al., 20206)
Applying such packages for TOME is slightly trickier, as TOME needs to perform ANNS inside the model. One viable route is to compute queries on-device, transmit them to a separate ANNS server and them transmit them back. We would recommend this approach for GPU accelerators, with faster host-device communication and slower device-device communication. As we are using TPU accelerators, decided to on-device ANNS, which does not require coordinating additional servers and will potentially allow for backpropagating through memory in future work.
C. ON-DEVICE NEAREST NEIGHBOR SEARCH
We shard the Mention Memory over all TPU devices. We perform search by distributing each query to all devices and retrieving top-K search results from each local memory shard. Then, the results are distributed back the original devices and the local search results aggregated through another, global top-K.
C.I.1 DOT-PRODUCT
The first method we describe is naive dot-product search, taking advantage of matrix multiplication capacity of TPU accelerators. In this method we perform search over local shards by taking the dot product between the query and the local memory shard and performing an approximate top-k operation the results. Dot-product search is easy to implement and fast for smaller memory sizes (up to 10 million entries). We implemented this method first due to its simplicity and our primary experimental results employ this search method.
14

---

Published as a conference paper at ICLR 2022
C.1.2 ANNS
To speed up search we implemented method 2) from standard CPU-based ANNS, bucketing the search table and searching only subset of buckets. In particular we perform k-means clustering to divide the Mention Memory into clusters, and perform dot-product search over the top ny clusters on each device.
C.1.3 OVERHEAD
While the Mention Memory is stored on-device, memory overhead negligible as the memory table sharded. For pre-training the Mention Memory took up 2.2% of available device memory. Table shows percentage of time spent on ANNS in TOME-1 pretraining for different reader architec- tures. The relative overhead of search becomes smaller with reader size, and ANNS overhead in particular becomes negligible for BERT-Large and up. We did not measure standard CPU ANNS overhead, but should be comparable or faster than our ANNS numbers.
Table 11: Proportion of time spent on ANNS TOME-1 pre-training setting.
Model
Dot-product
ANNS
BERT-Base
0.79
0.22
BERT-Large
0.48
0,07
T5-11B Encoder
0.17
0.02
C.1.4 HYPERPARAMETERS
For ANNS in TOMEBlocks we take top-2 search results from each local memory shard, and apply top- 128 over the retrieved results. For ANNS in the entity prediction layer take top-32 search results from each local shard, and aggregate across shards without applying an additional top-K operation.
D RETRIEVAL EXAMPLES
Table 12: TOME-2 retrievals for the second HoVer dev sample. We show top- retrieval results for the first (-+1) memory attention layer for passage mentions "the novel", "the movie" and "the album". Memory mentions in brackets. We non-named passage mentions, and generally understands it is looking for mentions related to music. However, while the best retrieval for "album" from a passage that mentions sampling the Shining, it is quite far removed and it is likely the retrieval is not sufficiently accurate here.
that the model retrieve relevant mentions for
Claim: Stephen King wrote the novel that the movie directed by Stanley Kubrick that was sampled in the album "Where Blood and Fire Bring Rest" based on. Label: TRUE
the novel -> Music Video. The video homage to Stanley Kubrick's 1980 film The Shining based on the [Stephen Kingl novel...
the movie -, Music Video. The video is a homage to Stanley Kubrick's 1980 film The Shining based on the [Stephen King] novel
the album ->1 Where Blood and Fire Bring Rest is the third full-length album released by [metalcore] band ZAO. It was the first album feature vocalist Dan Weyandt after the departure of Shawn Jonas along with new bassists/guitarists, Russ Cogdell and Brett Detar. The album contains sample from the film The Shining...
15


Reviews:
[{"review_id": "sRYKNOyoUrR", "reviewer": "Reviewer_giXh", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n- a batch-wise mechanism is proposed to support efficient pre-training.\n- entity-aware pre-training tasks are proposed to learn embeddings that can better embed entity knowledge.\n- evaluations conducted are comprehensive and sold.\n\nCons:\n- seems the proposed model still needs an NER to detect entity mentions from each input passage, which could lead to an unexcepted performance decrease when the model is used for out-domain inputs.\n- I suggest to include more entity-insensitive NLP tasks in the paper, to verify that proposed entity-aware mechanism is not biased on entity-related tasks.", "summary_of_the_paper": "This paper proposes an external knowledge-based pre-trained model that can leverage entity-wise knowledge embeddings. The mention encoder model encodes entity mentions occurred in an external corpus (i.e. Wikipedia) into embeddings as the memory. The TOME model retrieves most possible entity mentions from this memory and performs attention on them to generate aggregated embeddings, which will be integrated into the output representations of the current layer. To make training efficient, a two-stage training strategy is used where mention encoder is trained first and TOME is trained secondly. Evaluations are performed FEVER, HoVer, TriviaQA and ComplexWebQuestions, comparing with several baselines covering BERT, EaE, RAG and REALM, where TOME achieved consistent improvements on most of them.", "summary_of_the_review": "The paper is clearly written and easy to follow. The motivation of the proposed method is intuitive. The model design is reasonable. The experiments are solid. In general, this is a good paper. One suggestion is to add more entity-insensitive task, to verify that the proposed model will not decrease its performances on such tasks.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "nXtxkI9Txwq", "reviewer": "Reviewer_nJu8", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Pros\n- This work extends the past work on knowledge-augmented LMs and proposes an approach to encode knowledge sources in a more fine-grained way. \n- This architecture maintains explicit memories for mention spans in the knowledge source, so it still retains some interpretability.\n- Strong results on HoVer and CWQ. \n- The method and the training details are clearly explained.\n- The authors will release the code and model checkpoints.   \n\n### Cons \n- The model performance is inconsistent across different datasets. The TOME model shows strong results on HoVer but underperforms RAG on FEVER. Many FEVER claims are simple definitive sentences (e.g., single hop relations) and can be found in English Wikipedia, so it’s understandable that end-to-end retriever-based models might perform well. I wanted to see more analysis on this point (why is the number lower?). Similarly, it did well on CWQ, but the TriviaQA numbers are behind the retriever-based models. I think instance-level error analysis would be helpful to articulate underlying issues (from data or a model). \n- Although the authors promise to release their code and checkpoints, I’m wondering if people with limited resources (e.g., researchers in academia) can use this model. In Appendix B, it says that the TOME model is finetuned on 32 TPU. From the current version of the paper, it’s unclear if this model can be fine-tuned on a fewer number of TPUs (or GPUs?).\n\n### Questions / Suggestions\n- It seems that the pretraining data is annotated using off-the-shelf tools (NER, EL). I’m wondering about the trade-off between the entity coverage (e.g., adding more mention spans on top of the hyperlinks) and the data quality (cascading errors from the NER, EL tools). \n- If  the TOME model is only able to train on and answer 84% of examples, it might be interesting to see the baseline performance using the same train/eval sets. \n- There might be other choice for claim verification such as FM2 (https://arxiv.org/pdf/2104.04725.pdf) , which consists of higher quality claims compared to FEVER (but the data size is smaller).\n- The Fact as Experts paper is also related: https://arxiv.org/pdf/2007.00849.pdf\n\n### Minor\n- Use “English Wikipedia” in the main text (Bender rule).\n- Passage length is denoted as $T$ in Section 2 Notation, but $L$ is used in Appendix A. \n", "summary_of_the_paper": "__Method:__ This paper proposes a new approach to integrate knowledge sources into Transformer-based models. Concretely, entity mentions found in English Wikipedia (approx. 150M -- these mentions are linked to Wikipedia entities) are encoded into __Mention Memory__, which consists of key vectors and value vectors.  The knowledge representations in a __Mention Memory__ are accessed from a __TOME__ block, which is a stack of transformer blocks with a __memory attention__ layer. In this layer, a mention in the input sequence is converted into a “knowledge-injected” vector representation which is given by a weighted sum of the mention vectors (i.e., value vectors). The entire model (TOME) is pretrained in two stages. First, the Mention Encoder is pretrained, and the Mention Memory is generated. Here, the Mention Encoder is trained with a small-scale TOME architecture (Batch-TOME) using the masked LM loss and the coreference loss. Next, the Mention Memory is fixed, and other parameters in the TOME model are updated. The TOME model is trained with a masked LM task as well as an entity prediction task. This output layer allows us to use the TOME model for downstream tasks such as TriviaQA and Complex WebQA. \n\n__Evaluation:__ This approach is evaluated on two tasks: claim verification (HoVer, FEVER) and QA (TriviaQA, CWQ) and compared with various baselines (EaE, RAG, REALM etc.). In the claim verification experiments, the TOME models (both 1 and 2 blocks) outperform all baselines on HoVer, which require reasoning using multiple sources. On FEVER, the TOME models outperform baselines except RAG. In the QA experiments, the TOME model with 2 blocks outperforms all baselines on CWQ. On TriviaQA, the TOME models outperform similar models such as EaE and much larger generative LMs (T5) but underperform the retriever-based models. Additionally, the authors perform analysis on retrieved passages, memory sizes, and performance on unseen entities in a QA task.  \n", "summary_of_the_review": "I am leaning towards acceptance. The way TOME infuses textual knowledge into Transformer models is new. This approach actually integrates a passage retriever into a transformer-based LM and potentially covers broad downstream tasks. But, I think the analysis on the experimental results can be improved.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ekwWhsFRDna", "reviewer": "Reviewer_djR9", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n* While it's true that this approach combines aspects of previous work and is therefore not extremely novel, the combination is interesting and well motivated, and the paper is technically sound. This approach reminds me a lot of REALM and RAG, with the crucial difference that the retrieved knowledge isn't concatenated to the inputs but rather directly comes in the form of dense representations. This has a series of advantages, which the authors explain well in section 3 – no length constraint, less computationally expensive.\n* The paper is easy to follow. The appendices provide useful information on the pretraining procedure.\n* Performance on four standard tasks is convincing and the choice of baselines is reasonable.\n\nWeaknesses:\n* While pretraining is described very clearly, with full details of hyperparameters given in an appendix, details on the finetuning procedure are very light. \n* I would have liked to see a part of section 3 dedicated to discussing generative models like RAG.\n\nTypos:\n* p10: \"we have release\"", "summary_of_the_paper": "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.", "summary_of_the_review": "A strong, well written paper describing a novel approach for incorporating knowledge into a transformer encoder. The main components are not highly novel from a technical point of view, but their combination is.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "YmNVOwi4ZGe", "reviewer": "Reviewer_ShxU", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "STRENGTHS\n\n- The idea to attend over mentions rather than entities during feedforward is ambitious and refreshingly novel in the context of existing methods that focus on retrieving KB entries.\n\n- The paper develops a staged training scheme to make training feasible, which would be otherwise infeasible even with the compute resource the authors have.\n\n- TOME clearly outperforms EaE on the considered datasets with fewer parameters.\n\n\nWEAKNESSES\n\n- While TOME is certainly a neat extension of EaE, its main ideas are already pitched in previous works. Specifically, the benefits of virtual KBs (VKBs) over KBs, and performing multiple rounds of KB cross-attention within Transformer, are already highlighted in previous works. TOME's contribution is putting them together in a new model, but the main ideas individually are not a contribution of this work, so its novelty should be taken with a grain of salt.\n\n- While I understand that TOME is small and fast, and larger explicit retrieval-based models have an advantage over TOME, it substantially lags behind the existing state-of-the-art models on both fact checking and QA. I don't think the performance shouldn't be a reason to dismiss the contributions of the work, but it does mean that TOME is not state-of-the-art. \n\n- I think the paper can do better in justifying the choice of architecture and objectives. It's certainly reasonable, but there's no analysis. We don't know the impact of different span encoding architectures or pretraining objectives. Given how expensive it is to train TOME I suppose understanding development choices is a bit challenging. \n", "summary_of_the_paper": "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.", "summary_of_the_review": "The paper presents a large-scale pretrained Transformer that attends to mention encodings (EaE+VKB). It improves over EaE but lags behind state-of-the-art. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes TOME, which extends Transformer by attending to entity mention memory. Experiments are conducted on claim verification and QA.

Reviewers generally found the paper is solid. However, the novelty appears to be limited and is mainly in the combination of existing models.
Decision: Accept (Poster)