A Lightweight Information Extraction Model Based on DGCNN and Probabilistic Graphs
Background: A few months ago, Baidu held the "2019 Language and Intelligence Technology Competition", which had three tracks. I was quite interested in the "Information Extraction" track, so I signed up. After more than two months of grinding, the competition finally ended and the final results have been announced. Starting from knowing essentially nothing about information extraction, through the learning and research prompted by this competition, I eventually worked out some practical experience for doing supervised information extraction, which I'd like to share here.
Information Extraction track: Team "Scientific Spaces" ranked 7th on the final test results
I ranked 7th on the final test set, with an F1 of 0.8807 (Precision 0.8939, Recall 0.8679), about 0.01 behind the first-place team. That result isn't particularly outstanding from a competitive standpoint, but I believe the model has several genuinely novel aspects — such as a self-designed extraction structure, the use of CNN+Attention (hence quite fast), and the fact that it doesn't rely on pretrained models like Bert. I think this makes it a useful reference for both academic research and engineering applications in information extraction.
Basic Analysis
Information Extraction (IE) is a text processing technology that extracts factual information such as entities, attributes, relations, and events from natural language text. It is an important foundation for AI applications such as information retrieval, intelligent question answering, and intelligent dialogue, and has long attracted widespread attention in the field... This competition will provide the largest schema-based Chinese information extraction dataset in the industry (Schema based Knowledge Extraction, SKE), aiming to provide researchers with an academic exchange platform, further advance research on Chinese information extraction technology, and promote the development of related AI applications. — from the official competition website
Task Description
This information extraction task is, more precisely, a "triple" extraction task. Sample data looks like this:
{
"text": "Nine Mystic Beads is a novel serialized on Zongheng Chinese Web, written by Longma",
"spo_list": [
["Nine Mystic Beads", "serialization website", "Zongheng Chinese Web"],
["Nine Mystic Beads", "author", "Longma"]
]
}
The task is: given a sentence, output all the triples contained in it. Each triple has the form (s, p, o), where s is the subject — the primary entity, a span within the query — and o is the object — the secondary entity, also a span within the query — while p is the predicate, i.e., the relation between the two entities. The competition provided a candidate list of all predicates in advance (the schema, 50 candidate predicates in total). In short, (s, p, o) can be understood as "the p of s is o."
The competition provided nearly 200,000 annotated samples, with quite high annotation quality — thanks, Baidu. (Please don't ask me for the data; I'm not responsible for sharing the dataset. Apparently the dataset will eventually be released publicly at http://ai.baidu.com/broad/download, at which point you'll be able to download it.)
Characteristics of the Samples
This is clearly 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 as segmented by a tokenizer, so the query must be tagged directly in order to extract the correct s and o; and since word segmentation may get boundaries wrong, character-based input should be used for tagging.
2. Most samples yield extraction results of the form "one s, multiple (p, o)," e.g., "Wolf Warrior stars Wu Jing and Yu Nan," from which we should extract "(Wolf Warrior, star, Wu Jing)" and "(Wolf Warrior, star, Yu Nan)."
3. Samples with "multiple s, one (p, o)," or even "multiple s, multiple (p, o)," also make up a nontrivial fraction, e.g., "Wolf Warrior and Wolf Warrior 2 both star Wu Jing," from which we should extract "(Wolf Warrior, star, Wu Jing)" and "(Wolf Warrior 2, star, Wu Jing)."
4. The same (s, o) pair may correspond to multiple predicates p, e.g., "Wolf Warrior's star and director are both Wu Jing," from which we should extract "(Wolf Warrior, star, Wu Jing)" and "(Wolf Warrior, director, Wu Jing)."
5. In extreme cases, s and o may overlap. For example, "Lu Xun's Autobiography was published by Jiangsu Literature and Art Publishing House" should strictly yield not only "(Lu Xun's Autobiography, publisher, Jiangsu Literature and Art Publishing House)" but also "(Lu Xun's Autobiography, author, Lu Xun)."
Model Design
In the "Characteristics of the Samples" section we listed 5 basic observations. Aside from point 5, which is somewhat extreme, the other four are common characteristics of information extraction tasks in general. Before starting work in earnest, I briefly surveyed the mainstream information extraction models at the time, and found — to my surprise — that none of them covered all 5 characteristics well. So I abandoned the existing extraction paradigms and designed my own extraction scheme based on probabilistic-graph ideas, then, for the sake of efficiency, implemented it using CNN+Attention.
The Probabilistic-Graph Idea
For example, one fairly standard approach is to first perform entity recognition and then classify the relation between the recognized entities — but this approach cannot handle well the case where the same (s, o) pair corresponds to multiple predicates p, and it also suffers from sampling-efficiency issues. Another approach is to treat the whole thing as a single sequence-tagging problem, as in the paper Joint Extraction of Entities and Relations Based on a Novel Tagging Scheme — but this design cannot properly handle the case of multiple s's and multiple o's simultaneously, requiring an ugly "nearest-neighbor" heuristic. There's also the "using a sledgehammer to crack a nut" approach of resorting to reinforcement learning... And without exception, none of these methods can handle the case where s and o overlap.
Given how many years information extraction has been studied, it seems almost unbelievable to me that these basic issues have not been resolved. My own principle is: any inelegant design must be discarded. So I decided to abandon all the extraction paradigms I knew of and design my own extraction scheme. To that end, I borrowed a probabilistic-graph idea similar to seq2seq.
Anyone who has worked with seq2seq knows that the decoder is actually modeling
\begin{equation}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})\end{equation}
In actual prediction, we first use $x$ to predict the first word, then, assuming the first word is known, predict the second word, and so on recursively until an end token appears. So why not draw on this idea for triple extraction? Consider
\begin{equation}P(s, p, o) = P(s) P(o|s)P(p|s,o)\end{equation}
That is, we can first predict s, then feed s in to predict the o corresponding to that s, and then feed in both s and o to predict the relation p between the given s and o. In practice, we can merge the prediction of o and p into a single step, so the whole process needs only two steps: first predict s, then, given s, predict the corresponding o and p.
In theory, the above model can only extract a single triple. To handle the possibility of multiple s's, multiple o's, or even multiple p's, we use a "half-pointer / half-tagging" structure throughout (in plain terms, replacing softmax with sigmoid — this was also introduced in the earlier post A Reading-Comprehension-Style QA Model Based on CNN: DGCNN), and we also use sigmoid rather than softmax activation for relation classification.
With this design, the final model can be decoded very simply and efficiently, and it fully covers all 5 characteristics listed in "Characteristics of the Samples".
Note 1: Why not predict o first, and then predict s and the corresponding p?
That's because, in the second step, we need to sample and feed in the result of the first step (and sample only one). As already noted, in most samples the number of o's exceeds the number of s's, so if we predict s first and then feed s in to predict o and p, sampling over s is easy to make exhaustive (since there are few s's); conversely, if we had to sample over o it would be much harder to be exhaustive (since there could be many o's).
Keep this question in mind as you read on, and it will become clearer.
Note 2: While browsing recent arxiv papers, I found that, conceptually, this extraction design is similar to the approach in Entity-Relation Extraction as Multi-Turn Question Answering.
Overall Architecture
So far we have spent quite a lot of space explaining the extraction idea behind the model, namely: first identify s, then, given s, simultaneously identify p and o. Now let's look at the model's overall architecture.
To ensure efficiency, the model uses a CNN+Attention architecture (plus a short-sequence LSTM — since the sequence is very short, even an LSTM here doesn't hurt efficiency), and does not use notoriously slow pretrained models like Bert. The CNN follows the previously introduced DGCNN, and the Attention uses the Self Attention mechanism heavily promoted by Google. The overall architecture is shown below.
Diagram of this paper's information extraction model. Here the input sentence is "Wu Jing, the star of Wolf Warrior 2, was born in 1974", and the triples to be extracted are "(Wolf Warrior 2, star, Wu Jing)" and "(Wu Jing, date of birth, 1974)"
Specifically, the model's processing pipeline is:
1. Input a sequence of character IDs, and pass it through a hybrid character-word Embedding (the hybrid method is described later) to obtain the corresponding sequence of character vectors, then add Position Embeddings;
2. Feed the resulting "character-word-position embedding" into a 12-layer DGCNN for encoding, to obtain the encoded sequence (denoted $\boldsymbol{H}$);
3. Pass $\boldsymbol{H}$ through one layer of Self Attention, then concatenate the output with prior features (the prior features are optional; how they're constructed is discussed later);
4. Feed the concatenated result through a CNN and a Dense layer, and use a "half-pointer / half-tagging" structure to predict the start and end positions of s;
5. During training, randomly sample one annotated s (during prediction, iterate over all candidate s's), then feed the sub-sequence of $\boldsymbol{H}$ corresponding to this s into a bidirectional LSTM to obtain an encoding vector for s, then add relative-position Position Embeddings to obtain a vector sequence the same length as the input sequence;
6. Pass $\boldsymbol{H}$ through another layer of Self Attention, then concatenate the output with the vector sequence from step 5 and with prior features (again optional, described later);
7. Feed the concatenated result through a CNN and a Dense layer, and for each predicate p, build a "half-pointer / half-tagging" structure to predict the start and end position of the corresponding o — this way both o and p are predicted simultaneously.
This model differs considerably from an early baseline model I open-sourced (https://github.com/bojone/kg-2019-baseline), so please be aware of that.
Also, let me preemptively address two questions readers might have. First, "why sample only one s in step 5?" The answer is simple: one is enough (sampling more is equivalent to increasing the batch size), and sampling one is also easier to implement — I'd encourage readers who don't follow this to think it through carefully before continuing. Second, "why not use Bert?" This question is honestly a bit tedious — why do people keep asking it, can't I just not want to use it?... The real reason is that I've never been particularly fond of Bert, so I hadn't spent much time on fine-tuning it, until not long ago when I finally got hands-on with Bert fine-tuning — too late for this competition, though. Besides, fine-tuning based on Bert really isn't that interesting, it's inefficient, and it doesn't showcase much personal value-add, so unless it's really necessary, I'd rather not use it. (I did try a Bert-based approach a few days before the deadline; I'll write a separate post about that later.)
Model Details
Having covered the design philosophy and overall architecture, let's now look at the implementation details.
Hybrid Character-Word Embedding
As mentioned at the outset, to minimize boundary segmentation errors as much as possible, we should use character-based tagging, i.e., use characters as the basic input unit. However, plain character embeddings struggle to store effective semantic information — in other words, a single character carries essentially no semantics by itself. A more effective way of incorporating semantic information is a "hybrid character-word embedding."
In the model for this competition, I used a self-designed hybrid character-word scheme. First, we input the text as a character-level sequence and pass it through a character embedding layer to get a sequence of character vectors. We then segment the text into words and use a pretrained Word2Vec model to extract the corresponding word vectors. To align the word-vector sequence with the character-vector sequence, we repeat each word's vector as many times as the word has characters. Once we have this aligned word-vector sequence, we transform it via a matrix into the same dimensionality as the character vectors, and add the two together. The whole process is illustrated below:
Diagram of the hybrid character-word embedding scheme used in this model
In implementation, I used pyhanlp as the tokenizer, and trained a Word2Vec model (Skip-Gram + negative sampling) on 10 million Baidu Baike entries; the character vectors use a randomly initialized character embedding layer. During training, the Word2Vec word vectors are kept fixed, and only the transformation matrix and the character embeddings are optimized — which, from another angle, can also be seen as fine-tuning the Word2Vec word vectors via the character vectors and the transformation matrix. This way we incorporate the prior semantic information from a pretrained word-vector model while retaining the flexibility of character vectors.
By my own rough estimation, compared to using character vectors alone, this hybrid scheme improves the final results by about 1%–2%, which is a meaningful gain, and I've tested this scheme on other tasks too, obtaining similarly-sized improvements each time — confirming the effectiveness of this hybrid approach. Different pretrained word-vector models do have some effect on results, but not a huge one (less than 0.5%); I also tried using word vectors provided by Tencent AI Lab (using only the top 1 million words), and got roughly similar results.
Position Embedding
Since the model relies mainly on CNN+Attention for encoding, the resulting encoded sequence doesn't have a very strong "sense of position." But for this competition's data, positional information does carry real value — for instance, s usually appears near the beginning of the sentence, and o usually appears near s. One effective way to add positional information is via Position Embeddings, and unlike the formula-computed Position Embedding introduced previously, this model uses a learnable Position Embedding.
Specifically, we set a maximum length of 512 (as I recall, no sample sentence exceeds 300 characters), then initialize a new Embedding layer (with the same dimensionality as the character vectors) to all zeros. Given a position ID, it outputs the corresponding Position Embedding, which is added to the aforementioned hybrid character-word embedding to form the complete embedding, which is then fed into the DGCNN encoder described below.
Position Embeddings are also used elsewhere in the model, when encoding s: after the sampled s is encoded via a BiLSTM, we get a fixed-size vector, which we then replicate and concatenate into the original encoded sequence, serving as one of the conditioning inputs for predicting o and p. However, since o is more likely to be a word near s, rather than simply replicating the vector uniformly, I additionally add a "relative position vector," based on the current position relative to the position of s (if this description feels vague, please refer directly to the source code); this relative position vector shares the same Embedding layer as the one used for the initial input.
DGCNN
DGCNN was introduced previously in A Reading-Comprehension-Style QA Model Based on CNN: DGCNN — a design I proposed earlier when building a reading-comprehension model. It is essentially "dilated gated convolution." The concept of gated convolution comes from Convolutional Sequence to Sequence Learning, where it's called GLU (Gated Linear Units); I then replaced the ordinary convolutions with dilated convolutions to increase the receptive field. A similar approach appears in the paper Fast Reading Comprehension with ConvNets.
Combining residual connections with gated convolution to achieve multi-channel information transfer
When the input and output dimensions match, DGCNN can incorporate residual connections, and I previously showed that DGCNN with residuals is mathematically equivalent to a Highway-style dilated convolution:
\begin{equation}\begin{aligned}\boldsymbol{Y}=&\boldsymbol{X}\otimes \Big(1-\boldsymbol{\sigma}\Big) + \text{Conv1D}_1(\boldsymbol{X}) \otimes \boldsymbol{\sigma}\\ \boldsymbol{\sigma} =& \sigma\Big(\text{Conv1D}_2(\boldsymbol{X})\Big) \end{aligned}\end{equation}
The current model uses this form of DGCNN throughout, embodying selective multi-channel information transfer.
The final model uses a total of 12 DGCNN layers, with dilation rates of $[1, 2, 5, 1, 2, 5, 1, 2, 5, 1, 1, 1]$, i.e., $[1, 2, 5]$ repeated three times (learning repeatedly from finer to coarser granularity), followed by $[1, 1, 1]$ (fine-grained fine-tuning).
Distant-Supervision Prior Features
This competition did not allow the use of additional external triple knowledge bases, but we can assemble all the triples in the training set into a knowledge base, and then, when facing a new sentence, directly run a distant-supervision-style search against this knowledge base to obtain candidate triples for the sentence. "Distant supervision" here means: if two entities in a sentence happen to be the s and o of some triple in the knowledge base, that triple is extracted as a candidate triple. This way, given a knowledge base, we can use pure retrieval to obtain candidate triples for any sentence. Note, though, that these are merely candidates, and it's entirely possible that all of the extracted candidate triples are wrong.
My approach for using the distant-supervision results was: feed them into the model as features. First, all the s's obtained via distant supervision are assembled into a 0/1 vector structurally similar to the tagging scheme, which is concatenated onto the encoded vector sequence before predicting s. Then, all the o's and their corresponding p's obtained via distant supervision are likewise assembled into a similar 0/1 vector, concatenated onto the encoded vector sequence before predicting o and p. Please refer to the open-source code for implementation details. One caveat: during training, when constructing the distant-supervision features, the current training sample's own triples must first be excluded — i.e., only other samples' triples may be used to generate the distant-supervision result for the current sample — to properly simulate how it will be used at test time.
In terms of results: adding the distant-supervision prior features gave a very substantial improvement on the offline validation set — over 2%! And I've repeatedly confirmed there's no "future information leakage" bug in the code, so this offline result is legitimate. Unfortunately, though, the improvement on the online test set was about the same as without the prior features — essentially no gain. However, upon inspection, the outputs with and without the prior features differed quite a bit, so in the end I ensembled the two sets of results.
Other Supplementary Details
My implementation also includes some extra auxiliary modules, corresponding to the variables pn1, pn2, pc, po in the code. These modules theoretically provide some "global information": pn1 and pn2 can be thought of as a global entity-recognition module, pc as a global relation-detection module, and po as a global judgment of relation existence. None of these modules are trained separately; instead, their outputs are directly multiplied into the predictions for s and o.
These modules basically don't affect the final performance much, but they help speed up training, and intuitively I feel that including them is somewhat more principled.
Also, as mentioned earlier, after randomly sampling s, the vector sequence corresponding to s is fed into a BiLSTM for encoding. In practice this is implemented slightly differently: by uniformly interpolating between the start and end IDs of s, we extract a fixed number of vectors (6, in this model) to form a fixed-length vector sequence, which is then fed into the BiLSTM for encoding. This is mainly done to avoid dealing with variable-length s spans.
Experimental "Alchemy"
Finally, let's cover some details of the model training process.
Model code: https://github.com/bojone/kg-2019
The code was tested with Python 2.7 + Keras 2.2.4 + TensorFlow 1.8.
If this model's approach is helpful for your subsequent work, please consider giving credit (not that I expect much, really):
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.