LUKE : Deep Contextualized Entity Representations with Entity-aware Self-attention (2020)

Abstract: 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 .

## introduction many natural language tasks involve entities, e.g., relation classification, entity typing, named entity recognition (ner), and question answering (qa). key to solving such entity-related tasks is a model to learn the effective representations of entities. conventional entity representations assign each entity a fixed embedding vector that stores information regarding the entity in a knowledge base (kb) @xcite @xcite . although these models capture the rich information in the kb, they require entity linking to represent entities in a text, and cannot represent entities that do not exist in the kb. by contrast, contextualized word representations (cwrs) based on the transformer @xcite , such as bert @xcite , and roberta @xcite , provide effective general-purpose word representations trained with unsupervised pretraining tasks based on language modeling. many recent studies have solved entity-related tasks using the contextualized representations of entities computed based on cwrs @xcite @xcite . however, the architecture of cwrs is not well suited to representing entities for the following two reasons: (1) because cwrs do not output the span-level representations of entities, they typically need to learn how to compute such representations based on a downstream dataset that is typically small. (2) many entity-related tasks, e.g., relation classification and qa, involve reasoning about the relationships between entities. although the transformer can capture the complex relationships between words by relating them to each other multiple times using the self-attention mechanism @xcite , it is difficult to perform such reasoning between entities because many entities are split into multiple tokens in the model. furthermore, the word-based pretraining task of cwrs is not suitable for learning the representations of entities because predicting a masked word given other words in the entity, e.g., predicting "rings" given "the lord of the [mask]", is clearly easier than predicting the entire entity. in this paper, we propose new pretrained contextualized representations of words and entities by developing luke (language understanding with knowledge-based embeddings). luke is based on a transformer @xcite trained using a large amount of entity-annotated corpus question answering named entity recognition relation classification c 2 a beyoncé a [cls] c 1 a [mask] token emb. position emb. entity type emb. c 3 a in c 4 c 5 a los c 6 a [mask] bu beyoncé bu [mask] a [sep] c 7 d 2 transformer e e c 2 a beyoncé a [cls] c 1 bu [mask] a [sep] c 7 transformer with entity-aware self-attention e predict lives predict los_angeles pretraining entity typing predict location beyoncé [cls] [mask] in los [mask] beyoncé [mask] [sep] beyoncé [cls] [sep] ... ... ... ... [mask] h w1 h w2 h w3 h w4 h w5 h w6 h w7 h e1 h e2 h e1 h w7 h w2 h w1 words entities words entities d 5 + d 6 2 d 5 + d 6 2 predict angeles downstream tasks are solved using its output representations with linear classifiers. another key contribution of this paper is that it extends the transformer using our entity-aware self-attention mechanism. unlike existing cwrs, our model needs to deal with two types of tokens, i.e., words and entities. therefore, we assume that it is beneficial to enable the mechanism to easily determine the types of tokens. to this end, we enhance the self-attention mechanism by adopting different query mechanisms based on the attending token and the token attended to. we validate the effectiveness of our proposed model by conducting extensive experiments on five standard entity-related tasks: entity typing, relation classification, ner, cloze-style qa, and extractive qa. our model outperforms all baseline models, including roberta, in all experiments, and obtains state-of-the-art results on five tasks: entity typing on the open entity dataset @xcite , relation classification on the tacred dataset @xcite , ner on the conll-2003 dataset @xcite , clozestyle qa on the record dataset @xcite , and extractive qa on the squad 1.1 dataset @xcite . we publicize our source code and pretrained representations at https://github.com/studio-ousia/luke . the main contributions of this paper are summarized as follows: ## luke figure 1 shows the architecture of luke. the model adopts a multi-layer bidirectional transformer @xcite . it treats words and entities in the document as input tokens, and computes a representation for each token. formally, given a sequence consisting of m words w 1 , w 2 , ..., w m and n entities e 1 , e 2 , ..., e n , our model computes d-dimensional word representations h w 1 , h w 2 , ..., h wm , where h @xmath0, and entity representations h e 1 , h e 2 , ..., h en , where h @xmath1. the entities can be wikipedia entities (e.g., beyoncé in figure 1 ) or special entities (e.g., [mask]). ## experiments we conduct extensive experiments using five entityrelated tasks: entity typing, relation classification, ner, cloze-style qa, and extractive qa. we use similar model architectures for all tasks based on a simple linear classifier on top of the representations of words, entities, or both. unless otherwise specified, we create the input word sequence by inserting tokens of [cls] and [sep] into the original word sequence as the first and the last tokens, respectively. the input entity sequence is built using [mask] entities, special entities introduced for the task, or wikipedia entities. the token embedding of a task-specific special entity is initialized using that of the [mask] entity, and the query matrices of our entity-aware self-attention mechanism (q w2e , q e2w , and q e2e ) are initialized using the original query matrix q. ## analysis in this section, we provide a detailed analysis of luke by reporting three additional experiments. ## conclusions in this paper, we propose luke, new pretrained contextualized representations of words and entities based on the transformer. luke outputs the contextualized representations of words and entities using an improved transformer architecture with using a novel entity-aware self-attention mechanism. the experimental results prove its effectiveness on various entity-related tasks. future work involves applying luke to domain-specific tasks, such as those in biomedical and legal domains.