Efficient Entity Embedding Construction from Type Knowledge for BERT (2022)

Abstract: Recent work has shown advantages of incorporating knowledge graphs (KGs) into BERT for various NLP tasks. One common way is to feed entity embeddings as an additional input during pre-training. There are two limitations to such a method. First, to train the entity embeddings to include rich information of factual knowledge, it typically requires access to the entire KG. This is challenging for KGs with daily changes (e.g., Wikidata). Second, it requires a large scale pre-training corpus with entity annotations and high computational cost during pre-training. In this work, we efficiently construct entity embeddings only from the type knowledge, that does not require access to the entire KG. Although the entity embeddings contain only local information, they perform very well when combined with context. Furthermore, we show that our entity embeddings, constructed from BERT’s input embeddings, can be directly incorporated into the fine-tuning phase without requiring any specialized pre-training. In addition, these entity embeddings can also be constructed on the fly without requiring a large memory footprint to store them. Finally, we propose task-specific models that incorporate our entity embeddings for entity linking, entity typing, and relation classification. Experiments show that our models have comparable or superior performance to existing models while being more resource efficient.

## introduction many studies have attempted to enhance pretrained language models with knowledge such as ernie @xcite , knowbert @xcite , k-adapter @xcite ), e-bert @xcite , and kepler @xcite . among them, ernie, knowbert, e-bert, and kepler are typical work that do so by incorporating entity embeddings. the entity * work was extended after the internship with google. in this paper, we investigate whether we can construct entity embeddings by considering only local entity features. this is motivated by the observation that the context itself usually provides good information for the right answer. a number of examples are shown in one big advantage of utilizing only local features of entities is that we can update our entity embeddings very fast once the knowledge graph (kg) is changed, which is a desirable feature for kgs with rapid updates. we can construct the entity embeddings even on the fly to significantly reduce memory consumption and parameters since a number of tasks (e.g., entity linking) easily involve millions of entities. a disadvantage is that it is hard to infer the answer if large amounts of information are missing. for example, the language model analysis (lama) task @xcite re-quires a [mask] placeholder in the given sentence "sullivan was born in chippewa falls, wisconsin in @xcite " to be filled. the type knowledge may not be able to answer this question. thus, we do not focus on such tasks. instead, we apply our method on several typical entity-focused tasks, which were also chosen by related work. to construct the entity embeddings, we simply average bert's wordpiece embeddings from the type label of the entity as there are only 2.8 or 2.96 wordpiece tokens on average per label depending on our tasks. thus, our method is very fast and can be used to construct the entity embeddings on the fly without much cost to save memory and reduce parameters. for example, e-bert requires six hours to train its entity embeddings, while our method takes only about 1 minute to prepare the entity embeddings for our downstream tasks. the trained entity embeddings of e-bert take up around 30gb in size foot_0 . thus, storing these embeddings requires a large memory footprint, and the size continues to grow linearly if new entities are added. however, our method does not require such extra space for entity embeddings. for incorporation, previous work incorporates their entity embeddings during both fine-tuning and pre-training (ernie and knowbert). however, pre-training language models is a cumbersome and resource-intensive task. we show simply incorporating our entity embeddings during fine-tuning without any pre-training works well. one reason may be that these entity embeddings are directly constructed through averaging bert's wordpiece embeddings, so that they look like bert's word-piece embeddings, which may be helpful for incorporation for bert. finally, we propose task-specific models to incorporate our entity embeddings foot_1 . for entity linking, we propose a model that incorporates entity embeddings into the output; for entity typing and relation classification, the proposed model incorporates entity embeddings into the input. we show that our entity embeddings and incorporation method are simple and can achieve comparable or superior performance to existing methods on entity linking, entity typing, and relation classification. the contribution of this work can be summarized as follows: task example label entity linking cricket -england beat pakistan by 107 runs in second one-dayer. england_cricket_team pakistan_national_cricket_team entity typing gm is a publicly traded company that releases every bit of news they have organization relation classification cartí sugtupu can be reached by boat from the nearby onshore settlement of carti and the carti airstrip. place_served_by_transport_hub for the incorporation, ernie and knowbert both use new encoder layers to feed the entity embeddings, which requires pre-training. in contrast, e-bert achieves comparable results without pretraining by directly incorporating its entity embeddings into the standard bert model during taskspecific fine-tuning. one proposal from e-bert is to align the entity and bert wordpiece embeddings in the same space. to do so, it first trains word and entity embeddings jointly and then learns a linear mapping from word to bert wordpiece embeddings. the final entity embeddings can be obtained by applying this learned linear mapping so that they look like bert wordpiece embeddings. this mapping helps improve 4.4 micro f1 score on the test data on entity linking task. to learn this mapping, e-bert needs to train both word and entity embeddings, which are 30gb in size. our method for constructing entity embeddings shares the similar spirit, but it is an averaging method from bert wordpiece embeddings. k-adapter @xcite and kepler @xcite are both trained using multitask learning based on roberta @xcite in relation classification and knowledge base completion and do not rely on entity embeddings. outside the area of incorporating entity embedding into pretrained language model, there are a number of work that propose to use entity types from kgs on various tasks. for example, on entity linking task, some work use entity types together with entity descriptions or entity embedding trained over whole kg @xcite @xcite @xcite . some work use only entity types on entity linking task @xcite @xcite . @xcite use entity type embeddings for coreference resolution. the main difference between our work with them is that we mainly design our method for constructing entity embedding and our incorporation method for bert. as introduced before, we simply create entity embeddings from the bert's internal wordpiece embeddings. when incorporating our entity embeddings into bert, we also propose a model that makes use of bert's position embeddings on entity typing and relation classification task (mentioned in sec. 5.2). ## entity embedding construction we take the labels of the neighboring nodes for an entity obtained from wikidata as local features. since these labels are usually very short, as shown in figure 1 , we can efficiently obtain label embeddings by averaging wordpiece embeddings in the label. the final entity embeddings are then obtained by averaging the label embeddings. we denote m ij as the j-th wordpiece embeddings in the i-th label of an entity. the entity embeddings e are computed as follows: please note that m and n i are small in our relation classification task (1.27 and 2.96, respectively, on average). finally, the generated entity embeddings are updated in the task-specific fine-tuning. ## conclusion in this paper, we proposed to construct entity embeddings using local features instead of training