Triple Extraction with bert4keras

Model Introduction

The data format and the basic logic of the model were introduced in detail in the article "A Lightweight Information Extraction Model Based on DGCNN and Probabilistic Graphs", so I will not repeat them here. The dataset has been made public by Baidu and can be downloaded here.

Following the same strategy as before, the model is still based on the "semi-pointer, semi-annotation" approach for extraction. The sequence is to first extract the subject (s), and then pass s to extract the object (o) and predicate (p). The only difference is that the overall architecture of the model has been replaced with BERT:

1. Convert the original sequence into IDs and pass them into the BERT encoder to obtain the encoded sequence;

2. Connect the encoded sequence to two binary classifiers to predict s;

3. Based on the input s, extract the encoding vectors corresponding to the start and end of s from the encoded sequence;

4. Use the encoding vector of s as a condition to perform a Conditional Layer Norm on the encoded sequence;

5. Use the sequence after Conditional Layer Norm to predict the o and p corresponding to that s.

Class Imbalance

It is easy to imagine that when using a "semi-pointer, semi-annotation" structure for entity extraction, one faces the problem of class imbalance. Generally speaking, target entity words are much fewer than non-target words, so label 1 will be much less frequent than label 0. Conventional methods for handling imbalance, such as focal loss or manual adjustment of class weights, can be used, but after applying these methods, it becomes difficult to determine the optimal threshold. Here, I used a method I consider quite appropriate: raising the probability value to the n-th power.

Specifically, if the original output is a probability value p representing the probability of class 1, I now change it to p^n. That is, I consider the probability of class 1 to be p^n, while everything else remains the same. The loss function still uses the standard binary cross-entropy loss. Since 0 \leq p \leq 1, p^n will generally be closer to 0, making the initial state more consistent with the target distribution and thus accelerating convergence.

We can also compare the difference from the perspective of the loss function. Assuming the label is t \in \{0, 1\}, the original loss is: - t \log p - (1 - t) \log (1 - p) The loss after the n-th power becomes: - t \log p^n - (1 - t) \log (1 - p^n) Note that - t \log p^n = -nt \log p. Therefore, when the label is 1, it is equivalent to amplifying the weight of the loss. When the label is 0, (1 - p^n) is closer to 1, so the corresponding loss \log(1 - p^n) is smaller (and the gradient is also smaller). Thus, this can be seen as a way of adaptively adjusting loss weights (gradient weights).

Compared to focal loss or manual weight adjustment, the advantage of this method is that it makes the distribution closer to the target without changing the distribution of the original inner product (where p is usually obtained via an inner product plus a sigmoid), and not changing the inner product distribution is generally more friendly to optimization.

Source Code and Results

Github: task_relation_extraction.py

Without any pre-processing or post-processing, the final F1 score on the validation set is 0.822, which is generally better than the previous DGCNN model. Note that this is without any processing; if some pre- and post-processing were added, the F1 could likely reach 0.83.

At the same time, we can find that there are many errors and omissions in the annotations of the training and validation sets. When we participated in the competition, the annotation quality of the online test set was higher (more standardized and complete) than that of the training and validation sets. Therefore, the F1 submitted for testing was generally 4% to 5% higher than the offline validation F1. In other words, with some rule-based corrections, if this result were submitted to the leaderboard at that time, a single model would likely have an F1 of 0.87.

Points of Note

As mentioned at the beginning, I previously wrote an example using keras-bert to extract triples. Here, I will discuss the differences between the model in this article and the previous example, as well as some points worth noting.

The first difference is that the previous attempt was just a simple trial, so it only added the vector of s to the encoding sequence to predict o and p, rather than using Conditional Layer Norm as in this article. The Conditional Layer Norm approach has better representational power and yields a slight improvement in performance.

The second difference, which is also worth noting, is that the model in this article uses the standard BERT tokenizer, whereas the previous example used direct character-level segmentation. The sequences produced by the standard tokenizer are not simply split by character, especially when English and numbers are involved; the resulting tokens may not align with the characters of the original sequence. Therefore, extra care must be taken when constructing training samples and outputting results.

Readers might wonder: why not go back to the original character-level segmentation? I believe that since BERT is being used, one should follow the BERT tokenizer. Even if they don’t align, there are ways to handle it. The previous character-level segmentation was an unconventional usage caused by my lack of familiarity with BERT at the time, and it is not recommended. Following the BERT tokenizer may also result in better fine-tuning effects than forcing character-level segmentation.

Furthermore, I discovered a somewhat surprising fact: the vocabulary file (vocab.txt) provided with the Chinese BERT is incomplete. For example, the character "lu" (used in words like "fulu" or incantations) is not in BERT’s vocab.txt. Therefore, when outputting the final results, it is best not to use the tokenizer’s own decode method, but rather map back to the original sequence and use slicing on the original sequence for output.

Finally, the training of this version of the model includes an Exponential Moving Average (EMA) of weights, which can stabilize model training and even slightly improve performance. For an introduction to weight moving averages, you can refer to here.

Summary

This article provides an example of using bert4keras for triple extraction and points out several important considerations. Everyone is welcome to refer to it and try it out.