Recent advancements in natural language processing (NLP) have been driven by the utilization of large language models like BERT. These models, pre-trained on extensive textual data, capture linguistic and relational knowledge. Therefore, cloze-style prompts, which involve filling in missing words in a sentence, can be used to solve knowledge-intensive NLP tasks with the help of a language model. The "Knowledge Base Construction from Pre-trained Language Models (LM-KBC 2023)" challenge aims to harness language models' potential for knowledge graph construction through prompts. In particular, contestants are challenged to infer the correct Wikidata ID of objects, given a prompt used to link subject, relation, and object. Automatically inferring the correct objects would help in reducing the need for an expensive manual graph population. Our proposed approach in Track 1 focuses on expanding BERT's vocabulary with a task-specific one (i.e., Wikipedia2Vec) and facilitating its usage through prompt tuning with OPTIPROMPT.
Recently, significant progress has been made in enhancing downstream NLP tasks by leveraging
pre-trained language models like BERT [
In this regard, the challenge of Knowledge Base Construction from Pre-trained Language
Models (LM-KBC 2023)1 [
With the approach proposed in this work, we addressed Track 1 of the challenge and ranked second. The code for our suggested approach can be accessed on our GitHub repository2.
The goal of the knowledge graph construction task involves predicting the object entities
associated with a given subject entity-relation pair. For instance, if presented with the subject
entity-relation pair <Cologne - CityLocatedAtRiver>, the task is to predict the object entity
<Rhine>. In this challenge, when provided with an input subject-entity () and a relation (),
the objective is to employ language model probing to retrieve the complete set of accurate
object-entities (1, 2, ..., ) to solve the knowledge graph construction tasks. Language model
probing means asking the model to fill the masked token in a cloze-style prompt such as " Cologne
is a city situated along the [MASK] river." with the correct list of object entities, in this case
<Rhine>. Furthermore, the task requires linking this list of object entities to the correct Wikidata
[
The LM-KBC challenge provides a dataset with training and validation splits, along with the unseen test set. This data contains 21 distinct relations, for each of the relations one prompt template is given. Additionally, the dataset provides up to 100 training and validation examples per relation, each example containing a subject entity, the Wikidata ID of the subject entity, the relation, and a list of object entities with their respective Wikidata IDs. Four out of the 21 relations contain zero cases, where none of the object entities are correct. The dataset can be found at a Github3 repository, which also provides an in-depth dataset statistic.
For each example in the test set, Wikidata ID predictions are evaluated by calculating precision and recall against ground-truth values. The final macro-averaged F1-score is used to rank the participating systems submitted to the challenge.
The primary challenge with the task of using language models for knowledge graph construction
is that the language model is required to predict multi-tokens (e.g., <Pharrell Williams>) and
multi-label object entities (i.e., having 0 to N object entities for a given subject-relation pair).
While encoder-only language models like BERT [
As previously discussed in Section 1.4, using sequence-to-sequence models is not a viable option
for this task. We resort to using BERT since a preliminary experiment showed promising
results for predicting single token entities in particular. However, the main limitations of
BERT are the rather small and single-token-only vocabulary. To address this limitation, our
approach involves mapping BERT representations to a significantly larger and task-specific
representation space containing multi-token entities. One strategy for achieving this is by
utilizing word embeddings like Glove embeddings, comprising a 400K vocabulary size and
n-grams representation. However, GloVe embeddings [
Wikipedia2Vec [
In order to expand the vocabulary of BERT, we are required to map BERT word embeddings to
the Wikipedia2Vec embedding space. We accomplish this by pre-training a neural network (a
single-layer feed-forward linear mapping) denoted as : → 2 , which
transforms BERT embeddings to a corresponding vector representation in the Wikipeida2Vec
embedding space. To facilitate the training of this mapping as a first step, a dataset was
systematically generated. The generation process entailed string matching the token in BERT vocabulary
to their Wikipedia2Vec correspondents, while prioritizing entity embeddings over word
embeddings from Wikipedia2Vec. The resulting dataset denoted as = {(0, 0), (1, 1), ..., (, )}
consists of pairs (, ) where is a token embedding from BERT (e.g., the embedding of the
token ’France’ or ’hello’) and is the target embedding in Wikipedia2Vec space (e.g., the
embedding of the entity ’ENTITY/France’ or the word ’hello’, respectively). It’s important to note that
this simple string matching provides us with only positive examples. To generate hard negative
examples ˆ , we use a Faiss index [
exp− ((),)
Where is the Euclidean distance, ˆ are many nearest neighbors of in the Wikipedia2Vec embedding space and = 1000. This loss function incorporates principles from informa(1) tion theory, it not only separates positive and negative pairs but also maximizes the mutual information between () and the positive sample .
Using the available prompt templates6, we generated prompts for all subject-relation pairs
in the training set, provided by the organizers. For each prompt, we fetch the masked token
representation from BERT, denoted as ℎ, which is then transformed using the pre-trained
linear mapping to a representation in the Wikipedia2Vec space, as ℎ = (ℎ) (cf. Section
2.1). Subsequently, we search for the top 200 nearest neighbors of ℎ in the Wikipedia2Vec
space using a Faiss index, which will represent the negative examples ˆ = {ˆ0, ˆ1, ..., ˆ199}
in fine-tuning. To obtain positive examples = {0, 1, ..., } (where is the number of
correct object entity IDs for each subject-relation pair), we utilize the Wikidata object IDs
corresponding to each subject-relation pair to retrieve the respective Wikipedia articles and
subsequently their associated Wikipedia2Vec representations. For the zero object entity cases,
we incorporated the zero vector as the correct target. The primary objective of the fine-tuning
is to increase the Euclidean distance between the transformed masked token representations
ℎ and negative examples ˆ , while reducing the distance to positive examples . To achieve
this, the supervised contrastive loss (in Equation 2) as proposed by Khosla et al. [
1
∑︁ log ℒ = − | | ∈
exp(ℎ ,) ∑︀∈ exp(ℎ ,) + ∑︀^∈^ exp(ℎ ,^) = exp(− (ℎ , )) (2) (3)
To allow BERT to efectively adapt to the task, we perform prompt tuning using OPTIPROMPT
[
As depicted in Figure 1, for a given subject-relation pair, we first generate a prompt using the provided template and replace the relation tokens in the prompt with the optimal prompt tokens generated using OPTIPROMPT during the prompt and fine-tuning phase, as explained in Section 2.2 (cf. Step 1 in Figure 1). Following this, the modified prompt is employed to generate
the mask token representation using BERT (cf. Step 2 in Figure 1). This representation is then transformed using the fine-tuned linear mapping (cf. Step 3 in Figure 1). We utilize the Faiss index to fetch the top 200 closest (in terms of L2 norm) candidate entities to the transformed representation of the masked token. Simultaneously, we also calculate the similarity between candidate entities and the transformed masked token representation using Equation 1, which yields a probability distribution over the candidates (cf. Step 4 in Figure 1). To obtain the final set of object entities, we select the top answers where is inferred by = 1//() (cf. Step 5 in Figure 1). Finally, we use the Wikidata API to query the correct Wikidata IDs for the set of selected object entities (cf. Step 6 in Figure 1).
Detailed results of our approach on the test dataset are shown in Table 1. These outcomes provide insights into the performance of our approach across diferent relations, which is discussed below. In Table 2, a comparison is drawn between the results of our approach, the baseline method using BERT, and the winning approach.
The overall results of our approach are only marginally above the BERT baseline provided by the challenge organizers (cf. Table 2). The performance of BERT decreased in the relations concerned with countries and languages. One reason for this could be that BERT has the required words in its vocabulary. Thus, is able to produce meaningful outputs. On the other hand, we were able to increase the models’ performance for relations that depend on fine-tuning (i.e., <CompoundHasParts>, <PersonPlaysInstrument>, <PersonHasNobelPrize>). However, in this case, our model just learned to predict the most common object entities without much subject sensitivity. There are some relations that our model is now able to predict better, without a dataset-specific explanation. In particular, the increased performance for the < PersonHasAutobiography> and <CompanyHasParentOrganisation> show that we are indeed able to predict entities that were not originally in BERT vocabulary but learned from the Wikipedia2Vec embedding space.
In this work, a novel approach is introduced that aims to expand the BERT vocabulary by integrating entity-specific embeddings sourced from Wikipedia2Vec, in order to address the task of knowledge graph construction using language model probing. To summarize, our proposed approach comprises the following three steps: (i) Expanding BERT’s vocabulary by incorporating a larger and more task-specific representation through integration with the Wikipedia2Vec vocabulary. (ii) Training a mapping from the BERT embedding to the Wikipedia2Vec embedding space. (iii) Simultaneously performing prompt tuning using OPTIPROMPT and fine-tuning the mapping between BERT embedding to the Wikipedia2Vec embedding space. In conclusion, we observed that our model lacked sensitivity when predicting objects for diferent subject entities. We think that this behavior might be explained by three sources. Firstly, reducing the embedding size (768 (BERT) to 500 (Wikipedia2Vec)) thus losing parts of the information encoded in BERT. Secondly, Wikipedia2Vec knows more than 100 times the amount of words/entities than BERT. Therefore, we hypothesize that the embeddings might not carry enough information to distinguish between unknown representations. Lastly, it is not clear that BERT encodes the given subject texts in a meaningful way. Thus, we think that our pipeline would benefit from an end-to-end pre-training step, that increases both the sensitivity of the output representation and the knowledge available in the input representation. In the best case, retaining the contextual understanding of BERT while enabling the usage of a much larger and task-dependent vocabulary.