A Unified Encoder-Decoder Framework with Entity Memory (2022)

Abstract: Entities, as important carriers of real-world knowledge, play a key role in many NLP tasks.We focus on incorporating entity knowledge into an encoder-decoder framework for informative text generation. Existing approaches tried to index, retrieve, and read external documents as evidence, but they suffered from a large computational overhead. In this work, we propose an encoder-decoder framework with an entity memory, namely EDMem. The entity knowledge is stored in the memory as latent representations, and the memory is pre-trained on Wikipedia along with encoder-decoder parameters. To precisely generate entity names, we design three decoding methods to constrain entity generation by linking entities in the memory. EDMem is a unified framework that can be used on various entity-intensive question answering and generation tasks. Extensive experimental results show that EDMem outperforms both memory-based auto-encoder models and non-memory encoder-decoder models.

## introduction a large amount of real-world knowledge is related to entities, e.g., persons, nations, and events. entity knowledge is the information describing facts and attributes related to entities. many entity-intensive nlp tasks require models obtain entity knowledge to generate informative outputs, such as answering factual questions @xcite , explaining claims @xcite , or making informative conversations @xcite . pretrained encoder-decoder models can be directly applied on such entity-intensive tasks @xcite , but their ability to store and use knowledge is still questionable @xcite . a popular approach to incorporate knowledge into the generation process is retrieving evidence documents from external sources @xcite ; izacard and grave, transformer transformer who wrote [es] evening class [ee] ? + entity linking head entity memory transformer transformer [bos] [es] maeve binchy + [es] maeve binchy [ee] language modeling head encoder decoder 2021; @xcite . however, they suffer from significant computational overheads in indexing, retrieving, and reading a large number of extra documents @xcite . therefore, it is important to give encoder-decoder models access to entity knowledge without sacrificing too much efficiency. recently it has been proposed to use an in-model memory to augment auto-encoder models with entity knowledge on entity linking tasks @xcite @xcite . the entity memory stores entity knowledge as dense vectors which can be directly incorporated into the hidden states of transformer models @xcite , with no need to encode extra text. however, the auto-encoder framework in previous approaches can only select entities from a pre-defined entity vocabulary. hence, they are not able to give an entity outside the vocabulary, nor to generate answers or text beyond a single entity. in this paper, we propose a novel encoder-decoder framework with an entity memory (ed-mem), as shown in figure 1 . edmem is a unified framework on various entity-intensive qa and generation tasks, in which we train an entity memory for efficient knowledge incorporation. first, edmem is pre-trained on wikipedia documents, where it learns entity embeddings in the memory along with an encoder-decoder model. edmem learns to select relevant entities from the memory via an entity linking objective, and learns to generate answers using entity knowledge via a language modeling objective. second, to precisely generate entity names, we design three decoding methods that utilize the entity linking ability of edmem in its generation process, when we fine-tune it on downstream tasks. these include (1) free-form: left-to-right generation with entity identifiers; (2) static entity linking: first select entities by entity linking, build prefix trees for the selected entities, and then perform constrained entity generation using the trees; (3) dynamic entity linking: select entities on-the-fly for constrained entity generation. we conduct experiments on two popular testbeds of entity knowledge: open-domain qa and entityintensive generation. with the incorporation of entity knowledge, edmem outperforms non-memory encoder-decoder models on both tasks, and it retains the efficiency advantage of closed-book (i.e., non-retrieval) models. compared to memory-based auto-encoders, edmem achieves both higher overall accuracy (+9%) and better entity precision (+8%) on open-domain qa datasets, and it generates high-quality text from the memory-supported decoder on generation datasets when auto-encoders fail to do so. to summarize, edmem is the first knowledge-augmented closed-book framework to perform both tasks in a unified manner. open-book models open-book models first retrieve evidence documents from external corpora and read these documents to predict an answer @xcite . realm @xcite proposed a self-supervised approach to pretrain a retriever-reader model. dpr @xcite devised a contrastive objective to train a dense bi-encoder retriever on open-domain qa. subsequent approaches combined dpr with a generative objective to build large, powerful models on open-domain qa and generation tasks @xcite @xcite . however, openbook models have to process the raw text of all retrieved documents, which leads to extremely long inference time. besides, additional overheads are brought by loading the document index and retrieving evidence documents for each example. @xcite was the first to pre-train an entity memory with an autoencoder framework to perform entity prediction on open-domain qa. film @xcite followed eae and added a fact memory containing representations of wikidata triples. to better encode relational knowledge, opql @xcite learned latent relational representations for arbitrary entity pairs. recent work focused on learning a huge mention-level memory (~150m entries) with extensive pre-training @xcite or leveraging the entity memory in domain adaptive training @xcite . these models are all based on an auto-encoder framework. thus, they are able to predict entities ids but would fail to generate any non-entity answers or sentences. there is a preprint paper contemporaneous to our work which trained a memory with an encoder-decoder model @xcite . however, it used qa pairs as memory entries instead of entities, limiting its application to qa tasks. besides, their memory is much heavier (60m entries) than ours (1m). ## proposed framework suppose we have a pre-defined vocabulary of n entities @xmath0, . . . , e n }. a mention is the actual tokens in context which refer to an entity. the set of all mentions in the corpus is denoted as m. thus, there is a global alias table t : @xmath1, where each entity is mapped to all its mentions. the input of edmem is a sequence of tokens x of length s, and the target output is another sequence @xmath2, • • • , y t ] of length t . both sequences contain a pre-labeled set of mentions. each mention refers to an entity in e. we add two special tokens [e s ] and [e e ] to represent "entity start" and "entity end" boundaries of a mention, e.g., "[e s ] brett hart [e e ] is the president of the [e s ] united airlines [e e ]". these special tokens come from either wikipedia hyperlinks (in pre-training, §3.3) or an entity linking model (in fine-tuning, §3.4). ## experiments we test our edmem framework on two testbeds of entity knowledge: open-domain qa and entityintensive generation tasks. ## limitations first, if applying edmem to other datasets, its performance may correlate to the density of entity mentions in data examples. edmem may not be able to acquire sufficient entity knowledge from the memory if there are few mentions in the specific task. another limitation of our work is that the pre-trained entity memory may not be generalized to special domains, e.g., biomedical text. a lot of specific terminology is not included in our pre-trained entity memory, which may require additional training on domain-specific corpora.