A Lightweight Information Extraction Model Based on DGCNN and Probabilistic Graphs

Task Introduction

The information extraction task this time is, more precisely, a "triple" extraction task. Example data is as follows:

{
"text": "Jiu Xuan Zhu is a novel serialized on Zongheng Chinese Network, authored by Long Ma",
"spo_list": [
,

]
}

The goal is to input a sentence and output all triples contained in that sentence. A triple is in the form of (s, p, o), where s is the subject (a segment from the query), o is the object (also a segment from the query), and p is the predicate (the relationship between the two entities). The competition provided a list of all candidate predicates (schema, 50 in total). Generally, (s, p, o) can be understood as "the p of s is o."

The competition provided nearly 200,000 annotated data samples with high quality. (Please do not ask me for the data; I am not responsible for sharing the dataset. It is said that the dataset will be publicly released at http://ai.baidu.com/broad/download later.)

Sample Characteristics

Clearly, this is a "one-to-many" extraction + classification task. Through manual observation of the samples, the following characteristics were found:

1. s and o are not necessarily words identified by word segmentation tools. Therefore, the query must be labeled to extract the correct s and o. Considering that word segmentation might cut boundaries incorrectly, character-based input should be used for labeling.

2. Most extraction results in the samples are in the form of "one s, multiple (p, o)." For example, "The stars of ’Wolf Warrior’ include Wu Jing and Yu Nan," where we need to extract (Wolf Warrior, Starring, Wu Jing) and (Wolf Warrior, Starring, Yu Nan).

3. Samples with "multiple s, one (p, o)" or even "multiple s, multiple (p, o)" also account for a certain proportion. For example, "The stars of ’Wolf Warrior’ and ’Wolf Warrior 2’ are both Wu Jing," where we need to extract (Wolf Warrior, Starring, Wu Jing) and (Wolf Warrior 2, Starring, Wu Jing).

4. The same pair of (s, o) may correspond to multiple p. For example, "Wu Jing is both the star and director of ’Wolf Warrior’," where we need to extract (Wolf Warrior, Starring, Wu Jing) and (Wolf Warrior, Director, Wu Jing).

5. In extreme cases, s and o might overlap. For example, "’Autobiography of Lu Xun’ was published by Jiangsu Literature and Art Publishing House." Strictly speaking, besides (Autobiography of Lu Xun, Publisher, Jiangsu Literature and Art Publishing House), one should also extract (Autobiography of Lu Xun, Author, Lu Xun).

Probabilistic Graph Idea

For instance, a baseline approach is to perform entity recognition first and then classify relations for the identified entities. However, this approach cannot handle cases where one (s, o) pair corresponds to multiple p and suffers from sampling efficiency issues. Another approach is to treat it as a holistic sequence labeling task, referring to the paper "Joint Extraction of Entities and Relations Based on a Novel Tagging Scheme," but this design cannot handle multiple s and multiple o well, requiring ugly "proximity principles." There are also reinforcement learning methods... and without exception, none of these methods solve the problem of overlapping s and o.

It is quite incredible to me that after years of research in information extraction, these basic problems have not been solved. My principle is: unelegant designs must be discarded. Thus, I decided to abandon all extraction ideas I knew and design my own. For this, I considered a probabilistic graph idea similar to seq2seq.

Anyone who has worked with seq2seq knows that the decoder actually models: P(y_1,y_2,\dots,y_n|x)=P(y_1|x)P(y_2|x,y_1)\dots P(y_n|x,y_1,y_2,\dots,y_{n-1}) During actual prediction, the first word is predicted via x, then the second word is predicted assuming the first is known, and so on, until an end token appears. Why not apply this to triple extraction? We consider: P(s, p, o) = P(s) P(o|s)P(p|s,o) That is, we can first predict s, then pass s to predict the corresponding o, and then pass s and o to predict the relationship p. In practice, we can combine the prediction of o and p into one step. So the total process requires only two steps: first predict s, then pass s to predict the corresponding o and p.

Theoretically, the above model can only extract a single triple. To handle multiple s, multiple o, or even multiple p, we use a "half-pointer, half-labeling" structure (simply put, replacing softmax with sigmoid, as introduced in "A CNN-based Reading Comprehension Model: DGCNN"), and use sigmoid instead of softmax for relation classification.

With this design, the final model can decode very simply and efficiently, completely covering the five characteristics listed in "Sample Characteristics."

Note 1: Why not predict o first and then s and p?

Because in the second step of prediction, we need to sample and pass the result from the first step (and only sample one). As analyzed earlier, in most samples, the number of o is greater than the number of s. By predicting s first and then passing s to predict o and p, it is easier to sample s sufficiently (since there are fewer s). Conversely, sampling o would be harder to do sufficiently.

Note 2: Reading recent arXiv papers, I found that the idea of this extraction design is similar to the article "Entity-Relation Extraction as Multi-Turn Question Answering."

Overall Structure

We have spent considerable space explaining the extraction philosophy: identify s first, then pass s to identify p and o simultaneously. Now, let’s introduce the overall structure of the model.

To ensure efficiency, the model uses a CNN+Attention structure (plus a short-sequence LSTM; since the sequence is short, it doesn’t affect efficiency). It does not use pre-trained models like BERT, which are known for being slow. The CNN follows the previously introduced DGCNN, and the Attention uses Google’s Self-Attention. The overall structure is shown below.

Specifically, the model’s processing flow is:

1. Input the character ID sequence, obtain the corresponding character vector sequence through character-word hybrid embedding (details later), and add Position Embedding.

2. Input the "Character-Word-Position Embedding" into a 12-layer DGCNN for encoding to get the encoded sequence (denoted as \boldsymbol{H}).

3. Pass \boldsymbol{H} through a Self-Attention layer, then concatenate the output with prior features (optional).

4. Pass the concatenated result through CNN and Dense layers, using a "half-pointer, half-labeling" structure to predict the start and end positions of s.

5. During training, randomly sample one annotated s (during prediction, iterate through all s). Pass the sub-sequence of \boldsymbol{H} corresponding to this s into a BiLSTM to get the encoding vector of s. Add relative Position Embedding to get a vector sequence of the same length as the input sequence.

6. Pass \boldsymbol{H} through another Self-Attention layer, then concatenate the output with the vector sequence from step 5 and prior features.

7. Pass the concatenated result through CNN and Dense layers. For each p, construct a "half-pointer, half-labeling" structure to predict the start and end positions of the corresponding o. This predicts o and p simultaneously.

This model differs significantly from an earlier baseline model I open-sourced (https://github.com/bojone/kg-2019-baseline).

Regarding two possible doubts: First, "Why sample only one s in step 5?" The answer is simple: one is enough (sampling multiple is equivalent to increasing batch size), and one is easier to handle. Second, "Why not use BERT?" This is a bit of a moot point; I just didn’t feel like it. Specifically, I haven’t been a fan of BERT and hadn’t looked into fine-tuning it until recently, so I didn’t use it in the competition. Moreover, BERT fine-tuning is somewhat boring, inefficient, and doesn’t reflect personal value; unless necessary, I’d rather not use it. (I did try BERT a few days before the deadline; I’ll share that in another article.)

Character-Word Hybrid Embedding

As mentioned at the beginning, to avoid boundary errors, we use character-based labeling. However, pure character embeddings struggle to store effective semantic information. A more effective solution is "Character-Word Hybrid Embedding."

In this competition, I used a self-designed character-word hybrid method. First, we input the text sequence by character and get character vectors. Then, we segment the text and use a pre-trained Word2Vec model to extract word vectors. To align word vectors with character vectors, we repeat each word’s vector as many times as the number of characters in that word. After obtaining the aligned word vector sequence, we transform it to the same dimension as the character vectors using a matrix and add them together.

In implementation, I used pyhanlp for segmentation and trained a Word2Vec model (Skip Gram + Negative Sampling) on 10 million Baidu Baike entries. Character vectors used a randomly initialized embedding layer. During training, the Word2Vec vectors were fixed, and only the transformation matrix and character embeddings were optimized. This can be seen as fine-tuning Word2Vec vectors via character vectors and the matrix. This way, we integrate prior semantic information from pre-trained models while retaining the flexibility of character embeddings.

Based on my observations, this hybrid method improves the final result by about 1%–2% compared to pure character embeddings. This improvement is significant, and I have seen similar gains in other tasks. Different pre-trained models have some impact, but it’s usually small (less than 0.5%). I also tried using Tencent AI Lab’s word vectors (top 1 million words), with similar results.

Position Embedding

Since we mainly use CNN+Attention, the "sense of position" in the encoded vectors is not strong. However, for this competition’s data, position information is valuable (e.g., s often appears at the beginning, and o is usually near s). We used optimizable Position Embeddings instead of formula-based ones.

Specifically, we set a maximum length of 512, initialized a new zero-filled Embedding layer, and added the output to the hybrid embedding.

Another use of Position Embedding is when encoding s. The sampled s is encoded by BiLSTM into a fixed-size vector, which is then concatenated to the original sequence. To account for the fact that o is likely near s, I added a "relative position vector" (relative to s) to the copied vector, sharing the same Embedding layer as the input.

DGCNN

DGCNN was introduced in "A CNN-based Reading Comprehension Model: DGCNN". It is essentially "Dilated Gated Convolution." The gated convolution concept comes from "Convolutional Sequence to Sequence Learning", known as GLU (Gated Linear Units). I replaced standard convolutions with dilated ones to increase the receptive field. Similar approaches appear in "Fast Reading Comprehension with ConvNets".

When input and output dimensions are the same, DGCNN with residuals is mathematically equivalent to a Highway-style dilated convolution: \begin{aligned} \boldsymbol{Y} &= \boldsymbol{X} \otimes (1 - \boldsymbol{\sigma}) + \text{Conv1D}_1(\boldsymbol{X}) \otimes \boldsymbol{\sigma} \\ \boldsymbol{\sigma} &= \sigma(\text{Conv1D}_2(\boldsymbol{X})) \end{aligned} The model uses 12 layers of DGCNN with dilation rates [1, 2, 5, 1, 2, 5, 1, 2, 5, 1, 1, 1].

Prior Features from Distant Supervision

The competition prohibited external knowledge bases, but we can integrate all triples from the training set into a knowledge base and perform distant supervision searches for new sentences. If two entities in a sentence match an (s, o) pair in the knowledge base, that triple is extracted as a candidate.

I used these results as features passed into the model. Candidate s entities were formed into a 0/1 vector and concatenated to the encoding for s prediction; candidate o and p were similarly concatenated for o, p prediction. Note: during training, the current sample’s own triples must be excluded from the distant supervision feature to simulate test set conditions.

This feature improved the offline validation set by over 2%! However, online test results showed almost no improvement. Since the results with and without prior features differed significantly, I eventually fused them.

Other Supplementary Content

My implementation includes auxiliary modules (variables pn1, pn2, pc, po in the code) providing "global information." pn1 and pn2 are global entity recognition modules, pc is global relation detection, and po is global relation existence judgment. These are not trained separately but multiplied into the s and o predictions. They don’t significantly affect the final score but help speed up training.

Additionally, for encoding s, instead of the full BiLSTM, I used uniform interpolation between the start and end IDs of s to pick a fixed number of vectors (6 in the model) to form a fixed-length sequence for the BiLSTM, avoiding issues with variable-length s.

Basic Training Process

For the loss function, since "half-pointer, half-labeling" is essentially binary classification, I used binary cross-entropy. Note that s prediction has two binary classifications, while o prediction (including p) has 100 = 50 \times 2 binary classifications. Their losses are added 1:1. Thus, the absolute loss of o is 50 times that of s. This seems counter-intuitive, but experiments showed s and o are equally important. Cross-entropy variants like Focal Loss did not help.

I used the Adam optimizer, training with a 10^{-3} learning rate for up to 50 epochs, then 10^{-4} until optimal. The first epoch was used for WarmUp. I also used EMA (Exponential Moving Average) with a decay rate of 0.9999.

Decoding Optimization

The decoding process has room for tuning. For the "half-pointer, half-labeling" structure, I found that setting the "start" threshold to 0.5 and the "end" threshold to 0.4 yielded the best F1.

Furthermore, the test set was very high quality and standardized, while the training set had more omissions. I used post-processing rules to adjust predictions to meet official specifications, which provided nearly a 1% boost.

Model Averaging and Fusion

I used a layered ensemble scheme. Since models with and without prior features performed similarly but produced different results, I took their union. Before that, I used 8-fold cross-validation for each type (with and without prior features), resulting in 16 models. I averaged the probabilities of the 8 models in each group before decoding and then took the union of the two resulting files.

Knowledge Distillation

Since the test set was nearly perfect and the training set had omissions, I used a form of knowledge distillation to clean the training set. I used the 8 models (without distant supervision) to predict the training set. If a triple appeared in all 8 predictions but not in the original labels, I added it. If a triple appeared in none of the 8 predictions but was in the labels, I removed it. Retraining on this corrected set gave a nearly 1% boost.

Efficiency Strategies

"Efficiency" is a hallmark of this model. Even with a 16-model ensemble, predicting the 100,000 samples in the test set took only 4 hours (could be 2 hours with better CPUs). Other teams reported taking dozens of hours or even days.

For even more speed, one could reduce the 12-layer DGCNN to 6 layers, replace the BiLSTM for s with simple concatenation of start/end vectors, or use a smaller Word2Vec vocabulary. These could increase speed by over 5x with a 2%–4% performance drop.