BERT(S) for Relation Extraction in NLP
By: Wee Tee Soh
Tags: Naturallanguageprocessing, Entity Extraction, Relation Extraction, Bert, NLP

P re-training Task — Matching the Blanks (MTB)

How do you prepare an AI model to extract relations between textual entities, without giving it any specific labels (unsupervised)? Well, you will first have to frame the task/problem for the model to understand.

“Matching the blanks” between two relation statements. Snippet taken from https://arxiv.org/pdf/1906.03158.pdf

Consider the two relation statements above. Here, a relation statement refers to a sentence in which two entities have been identified for relation extraction/classification. Mathematically, we can represent a relation statement as follows:

Here, x is the tokenized sentence, with s1 and s2 being the spans of the two entities within that sentence. While the two relation statements r1 and r2 above consist of two different sentences, they both contain the same entity pair, which have been replaced with the “[BLANK]” symbol. Now, the intuition is that if both r1 and r2 contain the same entity pair (s1 and s2), they should have the same s1-s2 relation.

Therefore, the pre-training task for the AI model is that given any r1 and r2, to embed them such that their inner product is high when r1 and r2 both contain the same entity pair (s1 and s2), and low when their entity pairs are different. Noise-contrastive estimation is implemented here for this learning process, since it is not feasible to explicitly compare every single r1 and r2 pair during training.

The good thing about this is that you can pre-train it on just about any chunk of text, from your personal data in WhatsApp messages to open-source data on Wikipedia, as long as you use something like spaCy NER or dependency parsing tools to extract and annotate any two entities within each sentence.

Why the “[BLANK]” symbol then? Well, the entities within the relation statement are intentionally masked with “[BLANK]” symbol with a certain probability, so that during pre-training, the model can’t just rely on the entity names themselves to learn the relations (if it does that, the model will simply be memorizing, not actually learning anything useful), but also need to take into account their context (surrounding tokens) as well.

Pre-training Model — BERT

The model used here is the standard BERT architecture, with some slight modifications below to encode the input relation statements and to extract their pre-trained output representations for loss calculation & downstream fine-tuning tasks.

Entity Markers — Entity Start representation gives the best performance, according to the paper. Snapshot taken from https://arxiv.org/pdf/1906.03158.pdf

The above is what the paper calls Entity Markers — Entity Start (or EM) representation. In the input relation statement x, “[E1]” and “[E2]” markers are used to mark the positions of their respective entities so that BERT knows exactly which ones you are interested in. The output hidden states of BERT at the “[E1]” and “[E2]” token positions are concatenated as the final output representation of x, which is then used along with that from other relation statements for loss calculation, such that the output representations of two relation statements with the same entity pair should have a high inner product.

Once the BERT model has been pre-trained this way, its output representation of any x can then be used for any downstream task.

Downstream Task — Relation Classification

Suppose now we want to do relation classification i.e. given any two relations within a sentence, to classify the relationship between them (eg. Cause-Effect, Entity-Location, etc). Using the pre-trained BERT model on MTB task, we can do just that!

Adapted from https://arxiv.org/pdf/1906.03158.pdf

As above, simply stack a linear classifier on top of it (the output hidden states representation), and train this classifier on labelled relation statements. Thereafter, we can run inference on some sentences. The output, from me training it with the SemEval2010 Task 8 dataset, looks something like

In this case, the model successfully predicted that the entity “a sore throat” is caused by the act of “after eating the chicken”.

Downstream Task — One Shot Learning

Now, you might wonder if the model can still predict the relation classes well if it is only given one labelled relation statement per relation class for training. Well, it turns out that it can, or at least do much better than vanilla BERT models.

For the prediction, suppose we have 5 relation classes with each class only containing one labelled relation statement x, and we use this to predict the relation class of another unlabelled x. (Known as 5-way 1-shot) We can proceed to take this BERT model with EM representation (whether pre-trained with MTB or not), and run all the 6 x’s (5 labelled, 1 unlabelled) through this model to get their corresponding output representations.

We then simply compare the inner products between the unlabelled x’s output representation and that of all the other 5 labelled x’s, and take the relation class with the highest inner product as the final prediction.