PaperlayerView exact revision on arXiv

Source-linked AI summary

Entities as Experts: Sparse Memory Access with Entity Supervision

Thibault Févry, Livio Baldini Soares, Nicholas FitzGerald, Eunsol Choi, Tom Kwiatkowski

arXiv:2004.07202v2cs.CLcs.LG

TL;DR

The paper addresses how to capture declarative knowledge about entities in a language model’s learned parameters. It introduces EAE, which learns entity-specific memories from text, and reports strong knowledge-task performance while accessing only a small parameter proportion at inference.

  • Problem

    The paper seeks to develop neural sequence models that capture the knowledge required to answer questions about real-world entities.

  • Method

    EAE learns distinct entity representations from text and retrieves the correct memory for entity mentions during masked-span training.

  • Results

    EAE captures declarative knowledge across question answering, relation extraction, entity typing, and knowledge probing, outperforming related approaches with many more parameters.

  • Takeaways & Limitations

    Entity identification and representation are essential to EAE’s question-answering performance, while entity-specific parameterization enables sparse inference-time access.

  • Takeaways & Limitations

    EAE uses a pre-fixed entity vocabulary and cannot handle unseen entities; future work is needed for rare or unseen entities and less memory-intensive integration.

Abstract

from arXiv · show

We focus on the problem of capturing declarative knowledge about entities in the learned parameters of a language model. We introduce a new model - Entities as Experts (EAE) - that can access distinct memories of the entities mentioned in a piece of text. Unlike previous efforts to integrate entity knowledge into sequence models, EAE's entity representations are learned directly from text. We show that EAE's learned representations capture sufficient knowledge to answer TriviaQA questions such as "Which Dr. Who villain has been played by Roger Delgado, Anthony Ainley, Eric Roberts?", outperforming an encoder-generator Transformer model with 10x the parameters. According to the LAMA knowledge probes, EAE contains more factual knowledge than a similarly sized BERT, as well as previous approaches that integrate external sources of entity knowledge. Because EAE associates parameters with specific entities, it only needs to access a fraction of its parameters at inference time, and we show that the correct identification and representation of entities is essential to EAE's performance.

1 Introduction

EAE is a Transformer architecture that learns and retrieves entity-specific memories from text to capture knowledge about real-world entities. Its sparse entity-memory access supports question answering and other knowledge tasks while making correct entity identification central to performance.

  • Architecture: EAE learns distinct, independent entity representations directly from text and divides its parameter space according to entity identity.The model is named Entities as Experts because of this entity-specific parameter organization.
  • Architecture: Unlike a traditional Transformer, EAE can retrieve a dedicated memory for an entity across different textual mentions and reintegrate it into processing.The paper contrasts a dedicated Charles Darwin representation with ambiguity from token-based representations such as “Charles” and “Darwin.”
  • Training and inference: EAE is trained to predict masked spans while accessing memories only for entity mentions and retrieving the correct memory for each mention.Mention spans come from an existing detector, and entity identities come from Wikipedia hyperlinks.
  • Evaluation: EAE is evaluated on LAMA knowledge probes, open-domain TriviaQA, and WebQuestions, where it outperforms related approaches with many more parameters.The evaluation targets declarative knowledge and open-domain question answering.
  • Results: EAE’s learned entity representations outperform prior pre-trained embeddings on knowledge probing and TACRED relation extraction, while sparse access lets it use only a small parameter proportion at inference.The paper also reports that focusing the memory on entities is better than using an unconstrained memory store of similar size.

2 Approach

EAE interleaves a Transformer with an entity memory layer that retrieves sparse, entity-specific representations and trains mention detection, entity linking, and masked language modeling jointly.

  • Model Architecture: EAE represents each mention using a learned entity memory matrix and inserts retrieved representations sparsely into the Transformer sequence.The memory layer produces a projected weighted sum of entity embeddings at each mention’s start position.
  • Entity Memory: The model generates a pseudo entity embedding from mention-boundary representations, retrieves nearest entity embeddings, and forms a weighted memory representation.Training uses all entities for retrieval, while inference retrieves the top 100.
  • Model Architecture: The initial Transformer output predicts mention boundaries, retrieves entity embeddings, and supplies augmented input to subsequent Transformer layers.The final layer supports TokenPred and EntityPred heads.
  • Inference: EAE uses inferred BIO mention spans at inference rather than relying on an external mention detector.BIO labels encode Beginning, Inside, or Outside of a mention.
  • Learning Objective: Pre-training combines mention-boundary detection, entity linking, and masked language modeling losses using English Wikipedia mentions and hyperlinks.Entity linking supervision applies to hyperlinked mentions, while memory access for unlinked mentions is unsupervised.
  • Learning Objective: Entity-linking loss is absent for 88% of mentions without hyperlinks.The same loss is used for the EntityPred head, but memory access for those mentions remains unsupervised.

3 Related Work

EAE relates entity memory to knowledge-augmented language models, memory networks, and conditional computation while distinguishing its learned, entity-supervised access mechanism.

  • Knowledge-Augmented Language Models: Unlike approaches using external knowledge sources or precomputed entity representations, EAE learns entity representations from text with the model parameters.The related work contrasts EAE with KNOWBERT, which starts from BERT and integrates precomputed knowledge representations.
  • Memory Augmented Neural Networks: EAE’s entity memory can be viewed as a memory network whose slots correspond to learned entities and whose access is supervised through entity linking.This supervision is presented as a way to limit the computational cost of unconstrained memory access.
  • Conditional Computation: EAE applies conditional computation at two levels: it accesses memories only for entity mentions and retrieves only the needed entity memories.This routing is framed as treating entities as experts.

4 Models Evaluated

The evaluation compares EAE with ablations, BERT variants, retrieval-based QA systems, RELIC, and T5 across language modeling, entity linking, knowledge probing, and question answering.

  • EAE: EAE uses four Transformer layers before its entity memory layer and eight additional layers afterward, with TokenPred and EntityPred heads.Its Transformer-layer hyperparameters match BERT-base, and entity embeddings have size 256.
  • EAE: EAE retrieves k = 100 entity memories per mention at inference, enabling sparse activation of the entity embedding matrix.The paper notes that efficient top-k routing is not implemented in the current version.
  • Ablations: EAE-unsup removes entity-linking supervision from the memory while using full attention during inference.This ablation isolates the contribution of supervising memory slots with entity linking.
  • Ablations: The No-EAE baseline retains a similarly sized Transformer and prediction heads but lacks the entity-memory layer and its interactions between entity representations.Its entity embedding matrix is used only by EntityPred and cannot directly inform masked-language-model predictions.
  • Baselines: The comparison set also includes RELIC, T5 fine-tuned for open-domain question answering, and retrieval-based systems such as BM25+BERT, ORQA, and GraphRetriever.These models represent entity-embedding, encoder-decoder, and retriever-reader alternatives.

5 Knowledge Probing Tasks

EAE is evaluated on masked hyperlink prediction, LAMA knowledge probes, and open-domain question answering, with results showing benefits from supervised entity memory and entity-focused representations. Its advantages are strongest for entity-centered knowledge, while performance is limited for non-entity concepts, dates, and unresolved answers.

  • 5.1 Predicting Wikipedia Hyperlinks: EAE’s middle-layer entity memory improves both language modeling and entity linking over No-EAE, whereas unsupervised entity memory performs significantly worse.The results attribute this difference to the importance of entity-linking supervision for allocating the entity-memory parameters.
  • 5.1 Predicting Wikipedia Hyperlinks: EAE predicts mention tokens better than the 24-layer MM-large, but performs marginally worse in perplexity.The authors suggest that EAE may be overconfident when its token predictions are wrong.
  • 5.2 LAMA: EAE improves LAMA performance across entity-focused probes, while unsupervised entity memory performs worse overall.EAE’s average accuracy is similar to BERT-large, with complementary strengths across LAMA subtasks.
  • 5.2 LAMA: EAE is weaker on LAMA probes requiring non-entity concepts or dates, including ConceptNet answers such as “fly,” “cry,” and “happy.”The paper leaves modeling non-entity concepts and dates to future work.
  • 5.3 Open domain Question Answering: Entity prediction methods significantly outperform nearest-neighbor baselines, and adding an entity memory yields almost 6-point gains on both TriviaQA and WebQuestions.EAE is trained to predict linked answer entities rather than generate free-form text.
  • 5.3 Open domain Question Answering: Despite failing on 15% of TriviaQA and 9% of WebQuestions examples with unlinked answers, entity prediction outperforms standard T5 models, including one with 30× the parameters.T5-11B + SSM nevertheless surpasses EAE, motivating future combinations of entity prediction and text generation.

6 Analysis of TriviaQA Results

The analysis shows that EAE benefits from frequent and correctly linked entities, and can retain performance while retrieving far fewer entity embeddings. Its behavior also varies with question structure and entity coverage.

  • Entity-Based Analysis: EAE performs better on frequent entities, while performance is much lower for entities seen fewer than one hundred times.The authors attribute this likely to insufficient information in the corresponding entity embeddings.
  • Entity-Based Analysis: More named entity mentions improve EAE’s performance, whereas longer questions disadvantage closed-book models but not open-book models.The authors associate EAE’s improvement with access to more entity knowledge.
  • Manual Analysis: 87% of sampled TriviaQA questions contain no incorrectly predicted entities.The sample comprised 100 questions from the unfiltered-dev set.
  • Manual Analysis: Incorrectly linked entities considerably reduce EAE’s performance, even below examples with no named entities.This result was observed on the authors’ small manually analyzed sample.
  • Top-K over Entity Embeddings: Retrieving only K = 10 entity embeddings does not meaningfully deteriorate entity prediction or TriviaQA performance.The experiment varies the number of retrieved embeddings in the Entity Memory layer at inference time.
  • Top-K over Entity Embeddings: EAE accesses about 3% of the parameters while outperforming T5-3B on TriviaQA.The modular Entity Memory layer supports sparse parameter access, though naive implementation alone would not provide significant computational gains.

7 Comparison to Alternative Entity Representations

EAE learns entity representations jointly with a Transformer rather than relying only on fixed external representations. Comparisons indicate benefits from learned representations, while external initialization and knowledge sources remain useful avenues for improvement.

  • Learned Representations: EAE learns entity representations jointly with the parameters of its Transformer model.The comparison includes externally trained, fixed representations used by ERNIE and KNOWBERT.
  • Learned Representations: Learning entity representations is beneficial across the evaluated comparison settings.The authors compare random, TransE-Wikidata, and Deep-Ed initialization under frozen and jointly trained regimes.
  • Initialization Sources: Deep-Ed initialization improves most knowledge-probing tasks, while Wikidata initialization gives the best entity-typing results.The authors associate the entity-typing result with Wikidata’s coverage of person, organization, and location types.
  • Comparison Scope: The comparison is not isolated to entity representations because KNOWBERT and EAE differ in other architectural components.The authors specifically note KNOWBERT’s explicit entity-entity attention mechanism and additional WordNet synset embeddings.
  • Relation Extraction: EAE outperforms KNOWBERT on revised and weighted TACRED splits but slightly under-performs on the original setting.The result is reported for the cleaned TACRED dataset and the revised splits introduced by Alt et al. (2020).

8 Conclusion

The paper concludes that EAE is a Transformer architecture that learns entity representations from text and supports declarative-knowledge tasks. Its main scope boundary is that it cannot handle entities outside a fixed vocabulary.

  • Conclusion: EAE learns entity representations from text alongside the other parameters of a new Transformer architecture.The conclusion presents this learned entity-memory design as the paper’s central architectural contribution.
  • Conclusion: EAE captures declarative knowledge across open-domain question answering, relation extraction, entity typing, and knowledge probing.The conclusion lists these as representative task families supported by the evaluation.
  • Conclusion: The model’s entity representations influence open-domain question-answering predictions and are reported as high quality compared with KNOWBERT.This conclusion concerns the learned representations’ role and comparative quality.
  • Limitations: EAE cannot handle unseen entities because it learns representations for a pre-fixed entity vocabulary.The authors identify rare or unseen entities and knowledge-base integration as directions for future work.
  • System Design: EAE uses modular entity information rather than the standard open-book pipeline of document retrieval followed by a reader.Open-book systems retrieve documents or facts and feed them to a span-selection reader, whereas EAE organizes entity information in an Entity Memory layer.

B Wikipedia Pre-training

EAE is pre-trained from English Wikipedia contexts annotated with entity mentions and links. The corpus emphasizes frequent entities and uses additional unlinked spans for mention detection and retrieval.

  • Wikipedia Processing: The Wikipedia corpus contains 32 million contexts and over 17 million entity mentions.Contexts are created from the 2019-04-14 English Wikipedia dump by dividing articles into 500-byte chunks.
  • Wikipedia Processing: The training setup restricts the entity vocabulary to the one million most frequent entities, covering 86% of linked mentions.The data uses BERT tokenization with sequences limited to 128 word-piece tokens.
  • Wikipedia Processing: Unlinked mention spans are annotated with a mention detector and provide additional supervision for mention detection and retrieval.This allows retrieval even for mentions that are not linked in Wikipedia.
  • Pre-training: EAE is pre-trained from scratch for one million steps using ADAM, warmup, learning-rate decay, and gradient clipping.The setup uses a learning rate of 1e-4, warmup over the first 5% of training, and a clipping norm of 1.0.

C Open Domain Question Answering

The open-domain QA setup links proper-name mentions and eligible answers to entity IDs, while excluding unsupported training answers and reporting dataset-specific evaluation details. EAE is evaluated alongside related systems using prediction overlap and oracle accuracy.

  • Open Domain Question Answering: Answers in the entity vocabulary are linked to entity IDs with the SLING phrase table; unsupported answers are discarded from training but retained for development and test.This makes TriviaQA development-set oracle performance 85% because some answers are non-entities or outside the vocabulary.
  • Open Domain Question Answering: Questions are annotated with proper-name mentions using the paper’s mention detector, while nominal mentions are not linked because they may lack prior discourse introduction.This avoids forcing the model to hallucinate links for entities that have not been properly introduced.
  • Open Domain Question Answering: The open-domain QA setup reports linked-answer coverage for TriviaQA and WebQuestions, with the TriviaQA test-set figure unavailable because the test set is not public.The linked-answer coverage is summarized in Table 8.
  • Hyper-parameters: For TriviaQA and WebQuestions, models are selected through learning-rate and batch-size searches, with task-specific fine-tuning schedules and five seeds used for entity typing.The reported TriviaQA and WebQuestions schedules use 50,000 and 700 training steps, respectively.
  • Open Entity Processing: Entity typing uses nine generic types and Micro-F1 on 5,994 examples evenly divided among training, development, and test sets.Because mention boundaries and linking information are unavailable, entity embeddings and mention-detection parameters are frozen during fine-tuning.
  • Evaluation measures: Oracle accuracy counts a question as correct when either system answers correctly, while prediction overlap measures how often systems make the same predictions.Table 9 uses these measures to compare complementarity among QA systems.

E TriviaQA Evaluation Data Configuration

The TriviaQA evaluation follows prior open-domain conventions while aligning closely with T5’s setup. On TriviaQA Unfiltered-Dev, GR performs best, while oracle comparisons indicate complementary value across closed-book and open-book systems.

  • Dataset configuration: TriviaQA contains Wikipedia and Web-document blends, and open-domain approaches average over question-answer pairs for the Web data rather than using document-answer triplets.This distinguishes the open-domain evaluation from TriviaQA’s official web evaluation format.
  • Prior approaches: Prior approaches use different splits and reporting conventions, including web-data training and validation, Wikipedia test evaluation, or development-set reporting.These differences motivate describing the prior setup before comparison.
  • Our approach: EAE follows T5’s setup except that it does not retrain on combined training and development splits after hyper-parameter selection.This is the stated configuration for the closer comparison with T5.
  • Results: 55.4 accuracy for GR exceeds ORQA’s 45.1, EAE’s 43.2, and T5’s 42.3 on TriviaQA Unfiltered-Dev.The comparison includes two open-book systems, ORQA and GR, and two closed-book systems, T5 and EAE.
  • Complementarity: ORQA and GR have nearly 40% prediction overlap, while ORQA paired with T5 or EAE achieves higher oracle accuracy than ORQA paired with GR.The results suggest that some questions may be better suited to the closed-book paradigm, motivating future work on combining paradigms.

G Additional Examples of TriviaQA Predictions

Additional TriviaQA examples show that EAE benefits from correctly linked entities but remains vulnerable to date handling and entity-linking failures. Some errors are mitigated when other question entities are correctly identified.

  • Example organization: The examples are divided into questions with no named entities, correctly linked named entities, and incorrectly linked named entities.Table 10 presents questions, answers, and model predictions for the TriviaQA Unfiltered development set.
  • Questions without named entities: EAE fails on some date questions, whereas T5’s larger 11B-parameter model memorizes an obscure connection between ‘Sept 19’ and ‘talk like a pirate day’.In another entity-free example, EAE answers correctly when the answer entity is frequent in Wikipedia and connected to the question’s categorical information.
  • Correctly linked entities: Correctly linked entities enable EAE to answer one example correctly where T5 predicts a different director’s name.This example directly illustrates the usefulness of accurate entity linking for EAE predictions.
  • Date prediction: EAE also fails on date answers despite abundant question information, predicting the entity representing the 1990s.The example is presented as a failure in date prediction rather than a general failure to use question information.
  • Entity-linking failures: When an entity is unavailable, EAE may still answer correctly if two other question entities are predicted correctly; correcting a ‘Jimny’ typo enables the right entity link and answer.These cases show that partial entity identification can sometimes compensate for one unavailable or initially mislinked entity.
Loading 2004.07202v2…