Source-linked AI summary

LUKE: Deep Contextualized Entity Representations with Entity-aware Self-attention

Ikuya Yamada, Akari Asai, Hiroyuki Shindo, Hideaki Takeda, Yuji Matsumoto

arXiv:2010.01057v1cs.CLcs.LG

TL;DR

Entity-related tasks need effective contextualized representations of entities, but existing approaches rely on fixed knowledge-base embeddings or derive entity spans from word representations. LUKE treats words and entities as independent transformer tokens, pretrains on masked words and entities, and adds entity-aware self-attention. It achieves state-of-the-art results across five standard entity-related benchmarks.

  • Problem

    Existing entity representations either require entity linking and cannot cover out-of-knowledge-base entities, or do not directly produce span-level entity representations.

  • Method

    LUKE jointly represents words and entities as transformer tokens, pretraining on masked words and entities from Wikipedia and using entity-aware self-attention.

  • Results

    LUKE obtains state-of-the-art results on five tasks: entity typing, relation classification, NER, cloze-style QA, and extractive QA.

  • Takeaways & Limitations

    LUKE provides pretrained contextualized representations designed for entity-related tasks and directly models relationships between entities.

  • Takeaways & Limitations

    The reported pretraining uses original self-attention rather than entity-aware self-attention, whose query matrices are learned on downstream datasets.

Abstract

from arXiv · show

Entity representations are useful in natural language tasks involving entities. In this paper, we propose new pretrained contextualized representations of words and entities based on the bidirectional transformer. The proposed model treats words and entities in a given text as independent tokens, and outputs contextualized representations of them. Our model is trained using a new pretraining task based on the masked language model of BERT. The task involves predicting randomly masked words and entities in a large entity-annotated corpus retrieved from Wikipedia. We also propose an entity-aware self-attention mechanism that is an extension of the self-attention mechanism of the transformer, and considers the types of tokens (words or entities) when computing attention scores. The proposed model achieves impressive empirical performance on a wide range of entity-related tasks. In particular, it obtains state-of-the-art results on five well-known datasets: Open Entity (entity typing), TACRED (relation classification), CoNLL-2003 (named entity recognition), ReCoRD (cloze-style question answering), and SQuAD 1.1 (extractive question answering). Our source code and pretrained representations are available at https://github.com/studio-ousia/luke.

1 Introduction

LUKE introduces pretrained contextualized representations for words and entities, directly modeling both as tokens with entity-aware self-attention. It targets entity-related tasks and reports state-of-the-art results across five datasets.

  • LUKE addresses entity-related tasks by producing contextualized representations specifically for words and entities.Conventional entity embeddings require entity linking and cannot represent entities absent from the knowledge base, while contextualized word representations do not directly output span-level entity representations.
  • LUKE treats words and entities as independent transformer tokens, enabling direct modeling of relationships between entities.It computes intermediate and output representations for all tokens, including Wikipedia or special entities such as [MASK].
  • LUKE is pretrained by predicting randomly masked words and entities from a large entity-annotated Wikipedia corpus.Its entity-masking task extends BERT’s masked language model by replacing entities with [MASK] entities and predicting the originals.
  • Entity-aware self-attention distinguishes words and entities when computing attention scores.The mechanism uses different query mechanisms based on the attending token and the token attended to.
  • LUKE obtains state-of-the-art results on Open Entity, TACRED, CoNLL-2003, ReCoRD, and SQuAD 1.1.The evaluated tasks are entity typing, relation classification, named entity recognition, cloze-style question answering, and extractive question answering.

2 Related Work

Prior work includes static entity embeddings, contextualized word representations, and methods injecting knowledge into contextualized models. LUKE differs by learning entity representations during pretraining and using entity-aware self-attention.

  • Static entity representations use fixed embeddings learned from knowledge graphs or textual contexts associated with knowledge-base entities.These approaches store entity information independently of the surrounding text.
  • Contextualized word representations support entity-related tasks by computing entity representations from word representations.Representative systems include ELMo and BERT, based on bidirectional LSTM and transformer architectures.
  • Knowledge-enhanced contextualized models inject external knowledge using static entity embeddings, entity replacement objectives, knowledge-embedding objectives, or adapters.Examples include ERNIE, KnowBERT, WKLM, KEPLER, and K-Adapter.
  • LUKE differs by learning how to compute entity representations during pretraining and using entity-aware self-attention for entity-related tasks.The paper reports superior empirical results to existing contextualized and knowledge-enhanced contextualized representations in all experiments.

3 LUKE

LUKE represents words and entities as separate tokens in a bidirectional transformer, producing contextualized representations for both. It combines entity-aware self-attention with masked entity prediction during pretraining on Wikipedia.

  • Architecture: LUKE treats words and entities as input tokens and computes a contextualized representation for every token.Given m words and n entities, the model produces D-dimensional representations for both token types.
  • Input Representation: Token inputs combine token and position embeddings, while entity inputs additionally include a shared entity-type embedding.Entity names spanning multiple words use the average of the corresponding position embeddings.
  • Self-attention: Standard self-attention relates all m + n tokens through pairwise attention scores and weighted sums of transformed inputs.The query, key, and value matrices define the attention computation.
  • Entity-aware Self-attention: Entity-aware self-attention uses different query matrices for each pair of word and entity token types when computing attention scores.The mechanism distinguishes word-to-word, word-to-entity, entity-to-word, and entity-to-entity interactions without changing computational cost except during training.
  • Pretraining Task: LUKE pretrains with masked language modeling and masked entity prediction on a large Wikipedia corpus with hyperlink-based entity annotations.The final loss sums the MLM loss and cross-entropy loss for predicting masked entities from the entity vocabulary.
  • Implementation: The RoBERTaLARGE-based configuration uses 24 hidden layers, 16 attention heads, and approximately 483 million parameters.The model uses D = 1024 hidden dimensions, L = 64 attention-head dimensions, and H = 256 entity-embedding dimensions.

4 Experiments

LUKE is evaluated across five entity-related tasks using shared representation-based architectures and task-specific inputs. It outperforms strong baselines across the reported experiments, including state-of-the-art results on entity typing and relation classification.

  • 4 Experiments: LUKE is evaluated on entity typing, relation classification, NER, cloze-style QA, and extractive QA.The experiments use similar architectures with linear classifiers over word, entity, or combined representations.
  • 4.1 Entity Typing: Entity typing represents the target entity with a [MASK] entity and uses its representation for multi-label classification.The model is trained with binary cross-entropy loss averaged over entity types.
  • 4.1 Entity Typing: 2.0 F1 points over RoBERTa and 2.1 F1 points over KnowBERT: LUKE achieves a new state of the art on Open Entity.It also outperforms K-Adapter by 0.7 F1 points.
  • 4.2 Relation Classification: 1.4 F1 points over RoBERTa and 1.2 F1 points over MTB and KnowBERT: LUKE achieves a new state of the art on TACRED.It outperforms K-Adapter by 0.7 F1 points.
  • 4.3 Named Entity Recognition: 1.9 F1 points over RoBERTa and 0.8 F1 points over the previous state of the art: LUKE achieves a new state of the art on CoNLL-2003 NER.NER enumerates spans as entity candidates and classifies them into entity or non-entity types.
  • 4.4 Cloze-style Question Answering: 1.8 EM points and 1.9 F1 points over RoBERTa on ReCoRD development data: LUKE also surpasses RoBERTa ensemble on the test set without ensembling.ReCoRD selects the highest-scoring entity from passage candidates as the answer.

5 Analysis

The analysis tests whether LUKE’s entity representations and entity-aware self-attention drive its downstream performance, while controlling for longer pretraining. Ablations show both components contribute, and the advantage is not explained by training duration.

  • Removing entity inputs clearly degrades LUKE’s performance on CoNLL-2003 NER and SQuAD extractive QA.The ablation uses only word sequences while retaining the corresponding RoBERTa-style task architectures.
  • LUKE’s entity-aware self-attention consistently outperforms the original transformer attention across all evaluated tasks.The comparison is reported in the ablation study summarized in Table 7.
  • The largest improvements from entity-aware self-attention occur on TACRED relation classification and ReCoRD and SQuAD question answering.These tasks involve reasoning based on relationships between entities.
  • A RoBERTa model additionally pretrained for 200K steps achieves similar performance to original RoBERTa on CoNLL-2003 and SQuAD.This controls for whether LUKE’s gains result from its longer pretraining.
  • The additional-pretraining control indicates that LUKE’s superior performance is not owing to longer pretraining.

6 Conclusions

The paper concludes that LUKE provides contextualized representations for words and entities through an improved transformer with entity-aware self-attention. Its experiments demonstrate effectiveness across entity-related tasks, while future work targets biomedical and legal domains.

  • LUKE outputs contextualized representations of words and entities using an improved transformer architecture with entity-aware self-attention.
  • The experimental results demonstrate LUKE’s effectiveness on various entity-related tasks.
  • Future work includes applying LUKE to domain-specific tasks in biomedical and legal domains.
  • Table 7 reports an ablation study of LUKE’s entity-aware self-attention mechanism.
  • Table 8 reports results for RoBERTa additionally trained using the paper’s Wikipedia corpus.

A Details of Pretraining

The pretraining setup uses a large entity-annotated Wikipedia corpus and trains LUKE with specified hyperparameters and substantial computational resources. A separate RoBERTa continuation uses the same training duration and hyperparameters.

  • The pretraining corpus is the December 2018 Wikipedia release, containing approximately 3.5 billion words and 11 million entity annotations.Pages are split into sequences of at most 512 words with their hyperlink annotations.
  • LUKE pretraining uses AdamW with learning-rate warmup and linear learning-rate decay.The setup also updates only randomly initialized parameters to stabilize training.
  • The LUKE hyperparameters are documented in Table 9.
  • LUKE pretraining takes approximately 30 days on V100 GPUs.
  • The extra-pretrained RoBERTa model is trained for 200K steps and updates all parameters throughout training.Table 10 lists its hyperparameters.

B Details of Experiments

The experiments use constrained hyperparameter tuning, specified evaluation metrics, and documented common settings. SQuAD is evaluated without hyperparameter tuning using settings from prior work.

  • Experiments use NVIDIA’s PyTorch Docker container on a server with two CPUs and eight V100 GPUs.Hyperparameter tuning uses grid search on development-set performance for every dataset except SQuAD.
  • Performance is evaluated using EM on ReCoRD and F1 on the other datasets.
  • The constrained grid search tests 2, 3, and 5 training epochs.
  • SQuAD hyperparameters are not tuned and instead follow Liu et al. (2020).Other experiment details are reported in Tables 11 and 12.
  • Training uses AdamW with learning-rate warmup, linear decay, and early stopping based on development-set performance.

B.1 Open Entity

The evaluation uses five entity-related benchmarks with defined dataset sizes, label schemes, and official or code-based evaluation procedures. These include Open Entity, TACRED, CoNLL-2003, and ReCoRD in the supplied passages.

  • Open Entity contains 1,998 examples in each of its training, development, and test sets, labeled with nine general entity types.
  • TACRED contains 68,124 training, 22,631 development, and 15,509 test examples spanning 42 relation types.
  • Reported results use code based on Zhang et al. (2019) for Open Entity and TACRED, the conlleval script for CoNLL-2003, and official scripts or leaderboard submissions for ReCoRD.
  • CoNLL-2003 contains 14,987 training, 3,466 development, and 3,684 test sentences annotated for person, location, organization, and miscellaneous entities.
  • ReCoRD contains 100,730 training, 10,000 development, and 10,000 test questions based on 80,121 unique news articles.

B.5 SQuAD 1.1

SQuAD 1.1 contains question sets derived from Wikipedia articles, with development performance evaluated officially and test performance submitted to a leaderboard.

  • SQuAD 1.1 contains 87,599 training, 10,570 development, and 9,533 test questions based on 536 Wikipedia articles.
  • Development performance is computed with the official evaluation script, while test performance is obtained through leaderboard submission.

C Adding Entity Annotations to SQuAD dataset

Entity annotations are added to SQuAD by mapping names to Wikipedia referents and matching entity names in questions and passages. Ambiguous or low-link-probability names are excluded.

  • Entity names are mapped to referent Wikipedia entities using hyperlinks from each passage’s source Wikipedia page.
  • String matching extracts entity names from questions and passages, treating matched names as annotations for their referent entities.
  • Names referring to multiple entities are ignored, and names with link probability below 1% are excluded to reduce noise.
Loading 2010.01057v1…