Exploring Neural Entity Representations for Semantic Information

Andrew Runge
Duolingo
Pittsburgh, PA, USA
arunge@duolingo.com

Eduard Hovy
Carnegie Mellon University
Pittsburgh, PA, USA
hovy@cmu.edu

Abstract

Neural methods for embedding entities are typically extrinsically evaluated on downstream tasks and, more recently, intrinsically using probing tasks. Downstream task-based comparisons are often difficult to interpret due to differences in task structure, while probing task evaluations often look at only a few attributes and models. We address both of these issues by evaluating a diverse set of eight neural entity embedding methods on a set of simple probing tasks, demonstrating which methods are able to remember words used to describe entities, learn type, relationship and factual information, and identify how frequently an entity is mentioned. We also compare these methods in a unified framework on two entity linking tasks and discuss how they generalize to different model architectures and datasets.

Introduction

Neural methods for generating entity embeddings have become the dominant approach to representing entities, with embeddings learned through methods such as pretraining, task-based training, and encoding knowledge graphs (Yamada et al., 2016; Ling et al., 2020; Wang et al., 2019). These embeddings can be compared extrinsically by performance on a downstream task, such as entity linking (EL). However, performance depends on several factors, such as the architecture of the model they are used in and how the data is preprocessed, making direct comparison of the embeddings hard.

Some propose tasks that can only be applied to certain classes of embeddings, such as those produced from a mention of an entity in context. We address these gaps by comparing a wide range of entity embedding methods for semantic information using both probing tasks as well as downstream task performance. We propose a set of probing tasks derived simply from Wikipedia and DB-Pedia, which can be applied to any method that produces a single embedding per entity. We use these to compare eight entity embedding methods based on a diverse set of model architectures, learning objectives, and knowledge sources. We evaluate how these differences are reflected in performance on predicting information like entity types, relationships, and context words. We find that type information is extremely well encoded by most methods and that this can lead to inflated performance on other probing tasks. We propose a method to counteract this and show that it allows a more reliable estimate of the encoded information. Finally, we evaluate the embeddings on two EL tasks to directly compare their performance when used in different model architectures, identifying some that generalize well across multiple architectures and others that perform particularly well on one task.

1. Models

We compare eight different approaches to generating entity embeddings, organized along two dimensions: the training process of the underlying model, and the content used to inform the embeddings. Along the training dimension, the first method is task-learned embeddings, which are learned as part of a downstream task, such as EL. Pretrained embeddings are learned through a dedicated pretraining phase designed to produce entity embeddings. Finally, derived embeddings are produced by models capable of embedding any generic text, but that had no specific entity-based training.

Along the content dimension, the first type is description-based embeddings, which are learned or generated from a text description of the entity. Context-based embeddings are learned from words surrounding mentions of the entity. Lastly, graph-based embeddings are learned entirely from a knowledge graph, linking entities to types and to each other. Models may leverage multiple types of information to learn embeddings.

2.1 Task-Learned Embedding Models

First is the CNN-based model of Francis-Landau et al. (2016), a description and context-based hybrid model. It encodes text mentions of entities by applying convolutions over the mention’s name, context sentence, and the first 500 words of the document it appears in and encodes candidate KB entities with convolutions over the entity’s name and first 500 words of its Wikipedia page. It computes cosine distance between the outputs of each of the mention and KB convolutions, producing six features which are passed to a linear layer to produce a score for each candidate, trying to maximize the score of the true candidate. We use a kernel size of 150 and concatenate the candidate name and document convolution outputs to get 300-dimensional entity embeddings from this model.

Second is the RNN-based model of Eshel et al. (2017). The RNN outputs are each passed to an MLP attention module which uses the candidate entity embedding as the attention context to pass information from the text to the embeddings. The attention outputs and entity embedding are concatenated and passed through a single-layer MLP followed by a linear layer to compute a score for the candidate. We train these models using an EL dataset built from all of Wikipedia (Eshel et al., 2017; Gupta et al., 2017). We take the anchor text of each intra-Wiki link in Wikipedia as a mention, with the page it links to as the gold entity, filtering any cross-wiki links, non-entity pages, and entities with fewer than 20 words to create 93.8M training instances. Each mention is assigned a single negative candidate randomly from all entities (Eshel et al., 2017). We train each model for a single epoch on this dataset, following Eshel’s method.

2.2 Pretrained Entity Models

We evaluate three pretrained embedding models that leverage context and graph-based information to represent entities. For all three models, we train 300 dimensional entity embeddings. First is the context-based model of Ganea and Hofmann (2017) (Ganea). This model learns entity representations by sampling a distribution of context words around mentions of each entity and moves the entity embeddings closer to words in the entity’s context distribution and further from words sampled from a uniform distribution. The embeddings are normalized, resulting in a joint distribution of entities and words around the unit sphere, where entity vectors are close to their context words. We retrain their model on a larger subset of 1.5 million entities that includes the entities we use for the probing and EL tasks with a context window of 10 and 30 negative samples until it matches the authors’ original scores on an entity similarity metric (Ceccarelli et al., 2013).

We next use the graph-based BigGraph model of Yamada et al. (2018), a context and graph-based hybrid model which jointly trains word and entity embeddings (Wiki2V). It learns the word and entity embeddings using three tasks: 1) a skip-gram word prediction task, 2) an entity context task that predicts context words for each entity, and 3) an entity graph link prediction model that predicts which entities link to a given entity. We train the embeddings for 10 epochs, using the same context window and negative samples as the Ganea model.

2.3 Derived Models

Our first derived model is a simple bag of vectors model, in which we average the Google-News Word2Vec (Mikolov et al., 2013) vectors of the first 512 words of the entity’s Wikipedia page. Our other two derived models are BERT-based embeddings (Devlin et al., 2019) of the first 512 words in the entity’s Wikipedia page. We use BERT-base-uncased and BERT-large-uncased, which generate 768 and 1024 dimensional embeddings for each entity by averaging all the hidden states of all tokens in the final layer. We explored averaging the hidden states of the CLS tokens in different layers in initial experiments, but found averaging all hidden states in the final layer performed best.

3 Entity Embedding Probing Tasks

We next introduce a set of 22 probing tasks that can be applied to all of the embedding methods described above, divided into 5 categories based on the information they probe: context words used to describe a given entity, entity type information, relationships between entities, how frequently an entity is referenced, and factual knowledge.

Except for BigGraph, all of our models’ embeddings are trained on either text describing an entity or text surrounding mentions of an entity. As such, we explore how well the embeddings can recognize words used in the context of a given entity. We define an entity’s context words as the words which appear at least once in both 1) the first 500 words of the entity’s Wikipedia page and 2) a ten word window around an anchor link to that entity. By ensuring each word appears both in context with an entity and in the description, we can avoid biasing the task towards context or description-based embeddings. We create a binary prediction task for whether or not a word appears in an entity’s context words for 1,000 high frequency (appearing in >100k Wiki pages, W-H) and 1,000 mid-frequency words (>10k, W-M).

3.2 Entity Types and Sub-types

Similar to prior work (Yaghoobzadeh and Schutze, 2017; Chen et al., 2020), we examine how well different entity embedding methods are able to learn entity type information using probing tasks based on the DBPedia ontology. We extract the types from each of the first 3 levels of the ontology, representing increasingly fine-grained entity types, and create one N-way classification task for all types at that level, which we refer to as T-1, T-2, and T-3.

3.3 Relation Prediction

We probe for relationships between entities in three ways: 1) How reliably a relation type can be identified between a pair of entities, 2) how well the type of a relationship between a pair of entities can be predicted, and 3) how well the fact that two entities are related can be detected.

DBPedia only contains positive relationship examples between head and tail entities, so to create the binary tasks we must construct negative examples. We create negative examples by randomly replacing either the head or the tail in a positive example, weighted by how often the head entity appears as the head for this relationship type, and similarly for the tail entity (Wang et al., 2014). This reduces the risk of accidentally generating false negative corrupted relationships (true relationships that weren’t in DBPedia) by making it more likely that in N-to-1 or 1-to-N relationship types we replace the ‘1’ entity.

One risk with this approach is entity type leak solely on the two entities’ types, rather than true knowledge of their relationship. To address this, we modify the above replacement algorithm so that when replacing entity E, we select a replacement E′ that matches the entity type to the finest grained type possible in the DBPedia ontology. If the replaced entity has no type in the ontology, we select a random entity that has appeared in the same role (head or tail) for this relationship in the KB.

3.4 Entity Popularity

Prior work has found that incorporating the probability of an entity being linked to in a knowledge base is useful for downstream tasks such as EL (Yamada et al., 2016; Eshel et al., 2017). We define popularity as how frequently a given entity is linked to in Wikipedia from Wiki pages. We compute the popularity of each entity in Wikipedia and construct three types of tasks to probe for this information. First is a regression task, predicting the log-scaled number of times an entity is linked to (P-R). The second is a multi-class classification task for the binned number of times an entity is linked to as a coarser popularity estimate, with bins for > 1000, 100 – 1000, 10 – 100, and 1 – 10 links (P-B). The third is a comparative task, where the model must predict which of two entities is linked to more often. For fine-grained analysis, we select pairs for comparison based on the relative difference in their popularity. We create 3 tasks requiring one entity to have 2 (P-2), 5 (P-5) and 10 (P-10) times the number of links as its partner, and one unrestricted task (P-Any).

3.5 Factual Knowledge

Finally, we explore a small set of factual knowledge probes for spatial, temporal, and numeric information using triples of literals from DBPedia. The first two tasks probe if the embeddings retain the century or decade that a given person was born, based on the embedding for that person (F-C and F-D respectively). The next two tasks take as input a pair of location-type entities to see if the model can predict which of the two entities is larger in terms of 1) area in square kilometers (F-A) and 2) population (F-P). We select pairs using two methods, one which compares random pairs of entities and one that uses our type-restricted selection method from above to prevent the model from learning easy, type-based comparisons between, for instance, countries and villages, referring to the type-restricted versions as F-A+T and F-P+T. The final task compares two organisation type entities and tries to predict which has the higher revenue (F-R). We restrict pairs in this task to those whose revenues are reported in the same currency.

5 Downstream Task - Entity Linking

5.1 Experiments

Many of our embedding methods have been evaluated on EL tasks in prior work, either in a separate model or as full EL models themselves. However, direct comparison of the impact of the embeddings on EL performance is confounded by differences in the architectures which leverage the embeddings, as well as difficult to reproduce differences in candidate selection, data preprocessing, and other implementation details. To address this, we evaluate all of our embeddings in a consistent framework, testing them on two standard datasets in three different EL model architectures to directly compare the contribution of the embeddings to performance on the downstream task and how well they perform across different model architectures.