MGAD: Learning Descriptional Representation Distilled from Distributional Semantics for Unseen Entities

1 Introduction

Entity retrieval, which aims to efficiently filter a small number of candidate entities from a large knowledge base (KB) for a given mention, is a fundamental task for various natural language processing tasks, such as fact extraction and verification [Nooralahzadeh and Øvrelid, 2018]. Recently, embedding-based retrieval methods have achieved great success in entity retrieval due to their effectiveness in capturing the semantics [Wu et al., 2020b; Gillick et al., 2019; Botha et al., 2020]. In essence, these embedding-based methods represent entities and mentions with standalone encoders, in which entity embeddings can be pre-computed and stored offline for efficient retrieval.

A common paradigm for learning entity representations is to take attributes of each entity as input to the encoder, since these attributes consist of concise descriptions with rich semantic information. For example, [Wu et al., 2020b] purposes a BERT-based two-tower model to encode mentions and entity descriptions, and retrieves entities by nearestneighbor search. We denote these methods as descriptional representation methods. On the other hand, many recent studies propose to learn entity representations based on contexts around its mentions, which is grounded on the distributional hypothesis [Harris, 1954; Firth, 1957]. For example, LUKE [Yamada et al., 2020] takes anchor-texts in Wikipedia as mentions to learn entity embedding, then the embedding can be used in downstream tasks. We denote these methods as distributional representation methods. Distributional representation methods have achieved SOTA performances on several benchmark datasets. This is not surprising since distributional entity embeddings are more consistent with mention embeddings, making them clearly easier to match.

While distributional representation methods achieve better results than descriptional representation methods, the effectiveness of these methods highly relies on large amounts of occurrences for each entity. However, there are often lots of newly emerging entities, such as events and products, which are introduced continually along with the dynamic of KBs. For example, "COVID-19" is a new entity, which does not appear in the training corpus. These new entities often appear with only short descriptions accompanied with them, which we refer to as unseen entities. As a result, distributional representation methods cannot provide reliable representations for these unseen entities since there are no contexts for them in the training corpus.

In this paper, we aim to learn enhanced descriptional representations for unseen entities with inferred distributional semantics by employing a distillation framework. More precisely, we propose to leverage the distillation model to bridge the distributional embedding space and the descriptional embedding space, where the context-based encoder and the description-based encoder are considered as the teacher and the student, respectively. With the distillation model, we can transfer knowledge from the distributional embedding space (teacher) to the descriptional embedding space (student), and obtain the enhanced descriptional representations for unseen entities. To this end, we propose a novel distillation model referred as Multi-Granularity Alignments based Distillation (MGAD) model. Specifically, we first embed entity description into a vector space and then align the embedded entity from this space to a distributional semantic space. The distillation is implemented with four loss functions. Besides a retrieval loss, we propose three alignment losses with different granularities, i.e., element-level, batchlevel and space-level alignment loss. Experiments on four entity linking datasets show that MGAD performs competitively to the teacher model, and outperforms all other baselines. 1

2 RelatedWorks

In this section, we briefly review studies related to our work, including entity retrieval models, unseen entity representation learning, and lexical semantics alignment.

Entity Retrieval. Most entity retrieval systems consist of two stages, i.e., the candidate generation stage and the candidate ranking stage, to balance the efficiency and effectiveness [Wu et al., 2020b; Ganea and Hofmann, 2017; Onoe et al., 2021]. The candidate generation stage is to efficiently filter a small number of candidate entities from a large-scale KB, while the candidate ranking stage (also known as the disambiguation stage) chooses the most probable entity for each mention among the found candidates. In this paper, we focus on the candidate generation stage.

Researchers proposed different models for entity retrieval in the candidate generation stage. Traditional methods rely on heuristic functions [Le and Titov, 2019; Wu et al., 2020b; Yamada and Shindo, 2019], such as BM25 [Logeswaran et al., 2019] or alias table [Ganea and Hofmann, 2017] to build the model. Recently, neural models achieved great success in entity retrieval [Gillick et al., 2019; Botha et al., 2020; Wu et al., 2020a]. Existing neural models can be categorized into two classes, namely discriminative retrieval models and generative retrieval models. The former represents mentions and entities with dense vectors, and takes either simple interaction functions to evaluate their similarity [Gillick et al., 2019; Botha et al., 2020]. Example models like LUKE [Yamada et al., 2020] and BLINK [Wu et al., 2020b] take BERT or RoBERTa to encode contexts and descriptions to obtain entity embeddings, respectively. On the other hand, the latter retrieves entities by exploiting sequence-to-sequence architecture to generate entity names. For example, GENRE [Cao et al., 2021] takes a transformer architecture with a pre-trained language modeling objective to generate entity names.

1 Our data, code and models are available at https://github.com/ dalek-who/MGAD-entity-linking

UnseenEntity Problem. It is a unique problem for distributional representation methods in learning embeddings for unseen entities since they rely on contexts to learn the embedding. Instead, descriptional representation methods, e.g., BLINK [Wu et al., 2020b] and GENRE [Cao et al., 2021], are able to learn representations for unseen entities since they only rely on entity descriptions and titles. However, our experiments show that descriptional representation methods are less effective than distributional representation methods on seen entities. There are also some methods that employ external resources to learn embeddings for unseen entities. For example, DEEP-ED [Ganea and Hofmann, 2017] and CDTE [Gupta et al., 2017] learn embeddings from entity descriptions or types. ET4EL [Onoe et al., 2021] matches finegrained entity types without the need for entity embeddings. There are two tasks that are similar but different from our setting. The first is zero-shot entity linking [Logeswaran et al., 2019] which focuses on generalization on unseen domains. The second is unseen-mention entity linking [Onoe et al., 2021] which focuses on generalization on unseen mentions.

Lexical Semantics Alignment. Aligning embeddings from two different semantic spaces has been well studied in word representation learning [Wang et al., 2021; Zock and Schwab, 2008]. Researchers have proposed different models to learn word representations from descriptions, definitions, or contexts. For example, [Hill et al., 2016] encodes word definitions to fit the pretrained distributional word embedding by LSTM or Bag-of-Words model. [Bosc and Vincent, 2018] proposes to learn word embeddings based on dictionary definitions by reconstructing the word as well as the definition. [Bevilacqua et al., 2020] decodes a word with surrounding contexts into its definition by BART. Inspired by these studies, we try to address the distributional representation learning for unseen entities based on embedding alignment.

3 Our Approach

In this section, we will introduce our approach in learning enhanced descriptional representations for unseen entities by transferring knowledge from the distributional semantics based on the distillation framework. Specifically, our work is inspired by [Prokhorov et al., 2019], which employs canonical correlation analysis to align two embedding spaces, to learn semantic representations for unseen words. In this work, we borrow the idea of space alignment to learn embeddings for unseen entities. Instead of learning a linear mapping matrix, we propose a novel multi-granularity alignments-based distillation method to learn enhanced descriptional embedding for unseen entities. The overall framework is shown in Figure 1.

Formally, given a real-world knowledge base KB (e.g. Wikipedia), let Eall be the entity collection of KB. Each entity e is accompanied with a name namee and a description desce. Eall can be further divided into to two sub-collections of seen and unseen entities, namely Eseen and Eunseen. A seen entity (e ∈ Eseen) has contexts in KB besides its description, yet an unseen entity (e ∈ Eunseen) has only description without contexts.

3.1 Entity Embedding

In this subsection, we will introduce two different embedding methods to learn entity representations, namely descriptionbased and context-based entity embedding.

Description-based Embedding. The description-based embedding method takes only the entity description as input to learn entity embedding. In this work, we follow BLINK [Wu et al., 2020b] to take the pre-trained language model as the entity encoder. Specifically, for any entity e ∈ Eall with namee and description desce, the input Ie for entity encoder is constructed as:

Then the descriptional entity embedding ve is computed with:

where Ence(·) is a RoBERTa-like entity encoder.

Context-based Embedding. The context-based embedding methods learn distributional representations from entitycentric contexts for each entity. The state-of-the-art contextbased embedding method is LUKE [Yamada et al., 2020], which simultaneously learns contextualized representations of words and entities based on Transformer. In LUKE, for (and only for) any e ∈ Eseen, each entity embedding ue is randomly initialized and pre-trained with multitude of entitycentric contexts. This pre-training task is a variant of masked language model (MLM) task, where the model is trained to predict randomly masked entities in a context. In this work, we take LUKE as the teacher to guide our model to learn distributional representations for unseen entities.

The descriptional and distributional embedding spaces will be aligned by knowledge distillation, which will be introduced in Section 3.3. Before that, in the next section, we first introduce how mentions and entities are matched in retrieval.

3.2 Mention-entityMatching

In this subsection, we will introduce the mention encoder in LUKE and MGAD, and then introduce how to score a mention-entity pair (m, e).

MentionEncoder. Here, we directly employ themention encoder architecture of LUKE since we focus on the entity representation learning in this work. The mention encoder is a variant of RoBERTa and can encode several mentions simultaneously. Let Im be the input of n mentions with surrounding contexts:

Im = [CLS] c m1 c · · · mn c [SEP] [M]1 · · · [M]n, (3) where c is a non-mention word span, [M] is a special entitymask token and [M]i corresponds to the word span of the i-th mention mi by position embeddings. Then mention embeddings {vm1 , ..., vmn } are computed with:

where Encm(·) denotes the mention encoder, and {vm1 , ..., vmn } are the embeddings of [M]1 · · · [M]n in the last self-attention layer. It is worthy to note that this mention encoder benefits from long-distance contextual information and achieves SOTA performance on many entity-related tasks [Yamada et al., 2020].

Scoring Function. Given a mention-entity pair (m, e), ue ∈ RH denotes the entity embedding in LUKE, where H is the entity embedding size. ve ∈ RD denotes the descriptional entity embedding in MGAD, and vm ∈ RD denotes the mention embedding in both model, where D is the shared hidden size of Ence(·) and Encm(·) in Eq. (2) and (4). In MGAD, the matching score s(m, e)MGAD between a mentionentity pair is calculated by the dot-product similarity v⊤m · ve. In LUKE, due to the dimension miss-matching between vm and ue, a transformation ϕ is firstly applied on vm:

ϕ(vm) =W2 · LayerNorm(GELU(W1 · vm +b)) , (5) where ϕ(vm) ∈ RH . W1 ∈ RD×D , b ∈ RD and W2 ∈ RH×D are learnable parameters. Then the matching score s(m, e)LUKE can be calculated with dot-product ϕ(vm)⊤ · ue.

With the pre-computed and stored entity embeddings, give a mention embedding vm, both models can retrieve candidate entities by nearest-neighbor search.

3.3 Loss Function

Our distillation-based embedding alignment is implemented with four loss functions. Since we focus on entity retrieval, a retrieval loss Lre is necessary. Furthermore, there are three alignment losses with different alignment granularities, namely element-level alignment loss Lea, batch-level alignment loss Lba and space-level alignment loss Lsa. The final loss L of our distillation-based alignment is the weighted sum of the four losses:

Retrieval Loss. Retrieval loss Lre is for retrieval target itself. Following [Gillick et al., 2019], we use in-batch random negatives to sample positive and negative (m, e) pairs:

for each training batch, there are n randomly sampled mentions and their corresponding n entities (namely gold entities), and for each mention m, we construct one positive pair (m, e) with its gold entity e and n - 1 negative pairs (m, e) with the other entities e, which can be represented by an onehot label vector lbatch ∈ {0, 1}n . For mention m, let τ ∈ R be the temperature and sbatch m ∈ Rn be the scores over n entities, the score distribution στ (sbatch m ) for m over n entities can be defined as follows:

Let sbatch MGAD,m ∈ Rn be the score of mention m over n in-batch entities from MGAD and στ(sbatch MGAD,m ) ∈ Rn be the score distribution, then Lre can be defined with CrossEntropy:

Element-level Alignment Loss. As shown in Figure 2(a), element-level alignment loss Lea is an "one-to-one" alignment, where the descriptional embedding ve should be similar to the corresponding distributional embedding ue. In practice, because of the dimension miss-matching between ve and ue, a transformation function ψ(·) is first applied on ve where ψ(ve), ue ∈ RH . This is a usual solution for dimension missmatching in feature-based knowledge distillation [Gou et al., 2020]. Lea can be defined with the MSE loss:

Batch-level Alignment Loss. As shown in Figure 2(b), batch-level alignment loss Lba is an "one-to-batch" alignment, where the matching distribution of a mention over inbatch entities from the student is required to be similar to the teacher. Lba can be defined with KL-Divergence [Wu et al., 2020b]:

Lba = KL-Divergence (στ (sbatch LUKE,m ) , στ (sbatch MGAD,m )) , (10) where στ converts matching scores into probability distributions, sbatch LUKE,m and sbatch MGAD,m denote the matching scores of mention m from LUKE and MGAD. στ (sbatch MGAD,m ) ∈ Rn is the target distribution in KL-Divergence.

Space-level Alignment Loss. As shown in Figure 2(c), space-level alignment loss Lsa is an "one-to-all" alignment. In other words, one descriptional embedding ve should be able to select its corresponding ue from the entire distributional embedding space Sdist. Essentially, it can be seen as a "generation" task with vocab Eseen, whose label lvocab ∈ {0, 1}|Eseen| is a one-hot vector. For ve, Let svocab gen,e ∈ R|Eseen | be the "generation scores" over Eseen , στ (svocab gen,e ) ∈ R|Eseen | be the score distribution, svocab gen,e can be defined with dotproduct score:

Then Lsa can be defined with CrossEntropy:

4 Experiment

In this section, we conduct experiments to demonstrate the effectiveness of our proposed model on benchmark datasets.

4.1 Experiment Setup

We first introduce our experiment settings, including datasets, entity collections, baseline methods, training details, and evaluation metrics.

Datasets. To evaluate the performance of our model, we choose four widely used entity linking datasets: AIDA [Hoffart et al., 2011], ACE [Ratinov et al., 2011] , AQUAINT [Milne and Witten, 2008] and MSNBC [Cucerzan, 2007]. All documents in these datasets are manually annotated news articles. The collection of all gold entities in four datasets is denoted as Egold .

Entity Collections. As denoted in Section 3, following LUKE [Yamada et al., 2020], Eseen is the collection of the most popular 500K entites in Wikipedia, whose distributional embeddings are pretrained by LUKE. Furthermore, Eunseen is the collection of all e ∈ Egold - Eseen and 500K+ additional random entities. Eall = Eseen ∪Eunseen . Entity names and descriptions come from the Wikipedia dump cleaned by [Wu et al., 2020b].

Baseline Methods. We compare with five baselines:

• Alias: prior alias table is a heuristic method, in which candidate entities are retrieved by exactly alias matching and sorted by prior probability P (e|m). We directly take the implementation of REL [van Hulst et al., 2020].
• BM25: we follow [Wu et al., 2020b] to compare with BM25 where entities are indexed by titles.
• GENRE [Cao et al., 2021]: a BART-based state-of-theart generative model for entity retrieval.
• BLINK [Wu et al., 2020b]: a BERT-based two-tower retrieval model which learns description-based embeddings for entities.
• LUKE [Yamada et al., 2020]: a variant RoBERTa-based distributional two-tower model, which is our distillation teacher. Note that all experimental results of LUKE with Eunseen are empty because of the lack of distributional embeddings for Eunseen .

Training Details. Knowledge distillation aims to learn entity embeddings. Following [Yamada and Shindo, 2019], for further improving the matching ability, we infer the entity embeddings of Eseen, and finetune the mention encoder of MGAD on AIDA-train with entity embeddings fixed. For a fair comparison, other models are also finetuned in this way. Parameters are shared between Ence(·) in Eq. (2) and Encm(·) in Eq. (4). The initial parameters for Ence(·) and Encm(·) comes from a LUKE checkpoint, which is pretrained on Wikipedia but not finetuned on AIDA-train. ψ in Eq. (9) and (12) are the same with ϕ in Eq. (5). The best weights of four losses are α1 = 0.338, α2 = 0.002, α3 = 0.33, α4 = 0.33, where the sum of weights equals to 1. α2 is much smaller since the value of Lea is 102 times greater than others. Temperature τ = 1 in Eq. (8) and (12), and τ = 2 in Eq. (10). The optimizer is Adam [Kingma and Ba, 2015] with learning rate 2 × 10-5 and weight decay 0.01, linearwarmup learning rate scheduler with first 10% warmup steps. All models are trained to convergence with their best context lengths shown in Section 4.4.

Evaluation Metrics. We follow [Wu et al., 2020b] to report R@K (recall at K) where K=1, 10, 30. Moreover, we take MRR as an additional evaluation metric to evaluate the ranking performance. Since one document often contains multiple mentions, each metric is averaged in two strategies: 1) macro average for each document, and 2) micro average for each mention. We perform significant tests using the paired ttest. Differences are considered statistically significant when the p-value is lower than 0.05.

4.2 Overall Results

In this section, we compare MGAD against all baselines over seen entities and unseen entities separately. The main result of our experiments on the in-domain test set AIDA-testB are summarized in Table 1.

Firstly, we can see that the simple Alias model is a strong baseline that performs better than BM25. Recall that the Alias model works only on prepared alias names, which are often hard to obtain for unseen entities. Moreover, we find that neural methods (GENRE, BLINK, MGAD, and LUKE) generally perform better than heuristic methods (Alias and BM25) on top-ranked results (i.e., MRR and R@1), which indicates that entity retrieval can benefit from the semantic modeling ability of neural methods.

Secondly, we observe that LUKE performs consistently better than BLINK on all seen entities, which demonstrates the effectiveness of distributional representations over descriptional representations for seen entities. However, LUKE relies on large amounts of contexts to learn entity embeddings, making it unable to infer unseen entity embeddings.

Finally, we can see that our model MGAD achieves better performance than BLINK on almost all metrics over seen entities as well as unseen entities. Compared with LUKE, MGAD also obtains comparable performance on seen entities. Moreover, we can see that the relative improvement of MGAD over BLINK is greater in MRR and R@1 than in R@10 and R@30 on unseen entities, which shows the ability of MGAD to rank entities in top positions. All these results demonstrate that MGAD can learn better representations for entities on both seen entities and unseen entities.

4.3 Generalization Analysis

We further analyze the generalization ability of our model on 1) out-of-domain datasets and 2) entities with different frequencies. LUKE is excluded in the following experiments because of the lack of Eunseen embeddings.

Generalization On Out-of-domain Datasets. To analyze the generalization ability on out-of-domain datasets, we directly test MGAD on three benchmark datasets, i.e., ACE, MSNBC, and AQUAINT, without further finetuning. Results on these datasets are shown in Table 2. We can observe that: 1) the Alias model performs very well, which is consistently better than BM25. Moreover, it obtains relatively comparable performances to neural retrieval models, e.g., GENRE and BLINK. 2) Comparing GENRE and BLINK, we find that BLINK performs significantly better than GENRE when training and testing on AIDA dataset. While on three outof-domain test sets, the performance of BLINK drops significantly on ACE and AQUAINT, yet the performance of GENRE only drops slightly on ACE but improves significantly on MSNBC and AQUAINT. This demonstrates the robustness of the generative retrieval model. 3) On ACE and MSNBC, MGAD performs best among all baseline methods on all evaluation metrics. On AQUAINT, we observe that GENRE obtains good performance, this may be that entity names in AQUAINT are often copies or synonyms of mentions where GENRE can benefit from this bias.

Generalization On Entities With Different Frequencies. To further study the effectiveness of MGAD, we conduct experiments on AIDA to compare the performances on entities with different frequencies. Specifically, We split mentions into five bins according to the frequencies (on Wikipedia) of gold entities. Results are shown in Figure 3. We find some interesting trends where all three models perform relatively better on entities with much higher frequencies (e.g., frequency large than 104 ) or unseen entities. This is due to the fact that models can learn better representations for frequent entities. For those unseen entities, they are often distinct and show less ambiguity as they appear less frequent. This tells us that we should pay more attention to those entities with intermediate-frequency since they show the most ambiguity. Moreover, we can see that MGAD is consistently better than GENRE and BLINK on entities with different frequencies, this demonstrates the effectiveness of MGAD.

4.4 Ablation Study

To further analyze the impact of 1) four losses in distillation, and 2) context lengths in different mention encoders, we conduct ablation studies on AIDA dataset.

Loss Function. To study the impact of different loss functions, we drop one loss from L each time to see how the performance varies. All results are summarized in Table 3. Firstly, for seen entities, we can see that drop Lre impacts slightly but drop Lea and Lsa would significantly reduce the performance. Moreover, the performance reduces the most while dropping Lsa, which demonstrates the importance of the space-level alignment loss. Secondly, for unseen entities, we observe that the removal of each loss would reduce the performance significantly, and the impact of Lba is the greatest (although excluding Lba slightly benefits seen entities). The different effects of Lba on seen and unseen entities indicates that it prevents overfitting on seen entities. That is because that Lba aligns ve and ue through a mention, this indirect alignment can filter some unnecessary features. In contrast, Lea and Lsa directly align ve and ue, which is easy to be impacted by the noise in seen entities. Finally, although the overall performance with four losses is slightly poorer than excluding Lba , we still keep all losses in our model since we pay specific attention to unseen entities in this work.

Context Length. Since our model relies on mention contexts to learn embeddings for seen entities as the bridge between distributional embedding and descriptional embedding, we further study how the context length affects retrieval performances. To analyze the impact, we finetune mention encoders for GENRE, BLINK and MGAD with different context lengths. As shown in Figure 4, we can observe that with the context growing longer, MGAD and GENRE perform better, but BLINK performs worse. This indicates that the mention encoders of MGAD and GENRE can benefit from capturing long-distance contextual information, but that the mention encoder of BLINK lacks such ability and will be noised on the contrary. However, the worst performance of BLINK is still better than GENRE, which is another evidence of the effectiveness of two-tower models.

5 Conclusions and FutureWork

In this paper, we present a novel knowledge distillation framework to align descriptional entity embedding space with distributional entity embedding space, and take it to infer enhanced embeddings for unseen entities based only on descriptions. The experiments on four entity retrieval benchmarks demonstrate that our model can learn effective representations for unseen entities. For future work, we would like to analyze and address the ambiguity of intermediate-frequent entities in both entity retrieval and re-ranking stages.

Acknowledgments

This work was funded by the National Natural Science Foundation of China (NSFC) under Grants No. 61902381, No. 62141201, and 61872338, the Youth Innovation Promotion Association CAS under Grants No. 2021100, and 20144310, the Young Elite Scientist Sponsorship Program by CAST (No. YESS20200121), the Lenovo-CAS Joint Lab Youth Scientist Project, and the Foundation and Frontier Research Key Program of Chongqing Science and Technology Commission (No. cstc2017jcyjBX0059).