GPLinker: Joint Entity-Relation Extraction Based on GlobalPointer
Nearly three years ago, in Baidu's "2019 Language and Intelligence Technology Competition" (hereafter LIC2019), I proposed a new relation extraction model (see A Lightweight Information Extraction Model Based on DGCNN and Probabilistic Graphs), which was later further developed and published under the name "CasRel", and was considered SOTA for relation extraction at the time. However, when CasRel was proposed, I was actually encountering this field for the first time, so looking back now, CasRel still has quite a few imperfections. I've thought about further refining it since then, but never came up with a particularly good design.
Later, I proposed GlobalPointer and, more recently, Efficient GlobalPointer, and felt I now had enough "material" to build a new relation extraction model. So starting from the idea of probabilistic graphs and drawing on some of the SOTA designs that came after CasRel, I ended up with a model similar to TPLinker.
Basic Idea
At first glance, relation extraction looks like extracting triples $(s,p,o)$ (i.e., subject, predicate, object), but in terms of concrete implementation, it is actually the extraction of a "quintuple" $(s_h,s_t,p,o_h,o_t)$, where $s_h,s_t$ are the start and end positions of $s$ respectively, and $o_h,o_t$ are the start and end positions of $o$ respectively.
From the perspective of probabilistic graphs, we can build the model as follows:
1. Design a scoring function $S(s_h,s_t,p,o_h,o_t)$ for quintuples;
2. During training, make the annotated quintuples $S(s_h,s_t,p,o_h,o_t) > 0$, while all other quintuples are $S(s_h,s_t,p,o_h,o_t) < 0$;
3. During prediction, enumerate all possible quintuples and output the part that is $S(s_h,s_t,p,o_h,o_t) > 0$.
However, directly enumerating all quintuples would involve far too many candidates. Suppose the sentence length is $l$ and the total number of $p$ is $n$; then even with the constraints of $s_h\leq s_t$ and $o_h\leq o_t$, the total number of quintuples is still
\begin{equation}n\times \frac{l(l+1)}{2}\times \frac{l(l+1)}{2}=\frac{1}{4}nl^2(l+1)^2\end{equation}
This is a computational cost on the order of the fourth power of the length, which is infeasible in practice, so some simplification is necessary.
Simplified Decomposition
Given our current computational budget, we can generally afford at most a computational cost on the order of the square of the length, so at each step we can identify at most "one pair" of start or end positions. To this end, we can use the following decomposition:
\begin{equation}S(s_h,s_t,p,o_h,o_t) = S(s_h,s_t) + S(o_h,o_t) + S(s_h,o_h| p) + S(s_t, o_t| p)\label{eq:factor}\end{equation}
Note that this equation is a modeling assumption, designed based on our understanding of the task and the constraints of computational power, rather than something derived theoretically. Each term here has an intuitive meaning: for instance, $S(s_h,s_t)$ and $S(o_h,o_t)$ are the start/end scores for subject and object respectively, used to extract all subjects and objects via $S(s_h,s_t) > 0$ and $S(o_h,o_t) > 0$. The remaining two terms handle predicate matching: the term $S(s_h,o_h|p)$ represents matching using the start features of the subject and object as their own representations. If we can ensure that there is no nested entity within the subject or the object, then in theory $S(s_h,o_h|p) > 0$ alone would suffice to extract all predicates. However, considering the possibility of nested entities, we also need to perform an additional match on the end positions of the entities, which is the term $S(s_t, o_t|p)$.
At this point, the training and prediction process becomes:
1. During training, make the annotated quintuples $S(s_h,s_t) > 0$, $S(o_h,o_t) > 0$, $S(s_h,o_h| p) > 0$, $S(s_t, o_t| p) > 0$, while all other quintuples are $S(s_h,s_t) < 0$, $S(o_h,o_t) < 0$, $S(s_h,o_h| p) < 0$, $S(s_t, o_t| p) < 0$;
2. During prediction, enumerate all possible quintuples, successively output the parts satisfying $S(s_h,s_t) > 0$, $S(o_h,o_t) > 0$, $S(s_h,o_h| p) > 0$, $S(s_t, o_t| p) > 0$, and then take their intersection as the final output (i.e., all four conditions must be satisfied simultaneously).
In terms of implementation, since $S(s_h,s_t)$ and $S(o_h,o_t)$ are used to identify the entities corresponding to subjects and objects, this is equivalent to an NER task with two entity types, so we can accomplish this with a single GlobalPointer. As for $S(s_h,o_h| p)$, it is used to identify pairs of $(s_h,o_h)$ whose predicate is $p$. Unlike NER, it does not require the constraint $s_h \leq o_h$ here. We likewise use a GlobalPointer for this, but in order to identify the $s_h > o_h$ part, we need to remove the lower-triangular mask that GlobalPointer applies by default. Finally, $S(s_t, o_t|p)$ works the same way as $S(s_h,o_h| p)$, so we won't repeat the explanation.
Let's review this once more: as we know, as an NER module, GlobalPointer can uniformly identify both nested and non-nested entities, and it achieves this through token-pair-based recognition. So we should understand GlobalPointer more generally as a token-pair recognition model, rather than confining our understanding of it to the scope of NER. Once we recognize this, it becomes clear that the four scoring terms $S(s_h,s_t)$, $S(o_h,o_t)$, $S(s_h,o_h| p)$, $S(s_t, o_t|p)$ mentioned above can all be implemented with GlobalPointer, and whether or not to add the lower-triangular mask can simply be decided according to the specific task at hand.
Loss Function
Now that we've designed all the scoring functions, all that remains for training the model is the loss function. Here we continue to use the multi-label cross-entropy that GlobalPointer uses by default, as proposed in Generalizing "Softmax + Cross-Entropy" to Multi-Label Classification Problems. Its general form is:
\begin{equation}\log \left(1 + \sum\limits_{i\in \mathcal{P}} e^{-S_i}\right) + \log \left(1 + \sum\limits_{i\in \mathcal{N}} e^{S_i}\right)\label{eq:loss-1}\end{equation}
where $\mathcal{P},\mathcal{N}$ are the sets of positive and negative classes respectively. In previous posts, we always used "multi-hot" vectors to mark positive and negative classes: if the total number of classes is $K$, then we use a $K$-dimensional vector to represent them, where positions of positive classes are 1 and positions of negative classes are 0. However, in the scenarios of $S(s_h,o_h| p)$ and $S(s_t, o_t|p)$, we each need a $n\times l\times l$ matrix to represent the labels. Adding the two together and factoring in the batch size, the total dimensionality becomes $2bnl^2$. Taking $b=64,n=50,l=128$ as an example, we get $2bnl^2\approx 1\text{hundred million}$. This means that if we insist on representing the labels in "multi-hot" form, we would need to create a matrix with 100 million parameters at every training step and then transfer it to the GPU — both the creation and the transfer would be very costly.
So, to speed up training, we need to implement a "sparse" version of multi-label cross-entropy, where at each step we only pass the indices corresponding to the positive classes. Since positive classes are far fewer than negative ones, this greatly reduces the size of the label matrix. However, implementing a "sparse" version of multi-label cross-entropy means we need to compute equation $\eqref{eq:loss-1}$ while knowing only $\mathcal{P}$ and $\mathcal{A}=\mathcal{P}\cup\mathcal{N}$. To this end, the implementation we use is:
\begin{equation}\begin{aligned} &\,\log \left(1 + \sum\limits_{i\in \mathcal{N}} e^{S_i}\right) = \log \left(1 + \sum\limits_{i\in \mathcal{A}} e^{S_i} - \sum\limits_{i\in \mathcal{P}} e^{S_i}\right) \\ =&\, \log \left(1 + \sum\limits_{i\in \mathcal{A}} e^{S_i}\right) + \log \left(1 - \left(\sum\limits_{i\in \mathcal{P}} e^{S_i}\right)\Bigg/\left(1 + \sum\limits_{i\in \mathcal{A}} e^{S_i}\right)\right) \end{aligned}\end{equation}
If $a = \log \left(1 + \sum\limits_{i\in \mathcal{A}} e^{S_i}\right),b=\log \left(\sum\limits_{i\in \mathcal{P}} e^{S_i}\right)$, then this can be written as
\begin{equation}\log \left(1 + \sum\limits_{i\in \mathcal{N}} e^{S_i}\right) = a + \log\left(1 - e^{b - a}\right)\end{equation}
This way, we can compute the loss corresponding to the negative classes using $\mathcal{P}$ and $\mathcal{A}$, while the loss for the positive classes remains unchanged.
Finally, in general multi-label classification tasks, the number of positive classes is not fixed. In such cases, we can index classes starting from 1, using 0 as a padding label so that the label matrix for every sample has a consistent size, and then mask out class 0 in the loss computation. The corresponding implementation is already built into bert4keras — see "sparse_multilabel_categorical_crossentropy" for details.
Experimental Results
For convenience, let's refer to the model above as GPLinker (GlobalPointer-based Linking). A reference implementation based on bert4keras is provided here:
Script link: task_relation_extraction_gplinker.py
The experimental results on LIC2019 are as follows (CasRel's code is task_relation_extraction.py):
$$\begin{array}{c|c} \hline \text{model} & \text{F1} \\ \hline \text{CasRel} & 0.8220 \\ \text{GPLinker (Standard)} & 0.8272\\ \text{GPLinker (Efficient)} & 0.8268\\ \hline \end{array}$$
The pretrained model is BERT base; the difference between Standard and Efficient is that they use standard GlobalPointer and Efficient GlobalPointer respectively. These results demonstrate two things: first, that GPLinker is indeed more effective than CasRel; second, that the design of Efficient GlobalPointer can indeed match the performance of standard GlobalPointer while using far fewer parameters. It's worth noting that on the LIC2019 task, if standard GlobalPointer is used, GPLinker's parameter count is close to 10 million, whereas with Efficient GlobalPointer it is only about 300,000.
Additionally, on a 3090 GPU, compared with the "multi-hot" version of multi-label cross-entropy, the model using the sparse version achieves a 1.5x speedup in training without any loss of accuracy. Compared with CasRel, GPLinker with sparse multi-label cross-entropy is only 15% slower in training, but nearly twice as fast in decoding — making it both faster and better.
Related Work
For readers familiar with the progress of SOTA relation extraction models over the past couple of years, once you understand the model above, you'll notice it is quite similar to TPLinker. Indeed it is — the model was designed by drawing heavily on TPLinker, and the final results also turned out to be quite similar to TPLinker's.
Broadly speaking, the differences between TPLinker and GPLinker are as follows:
1. TPLinker's token-pair classification features are obtained by concatenating the start and end features and then applying a Dense transformation, an idea derived from Additive Attention; GPLinker instead uses GlobalPointer, an idea derived from Scaled Dot-Product Attention. On average, the latter has lower GPU memory usage and faster computation.
2. GPLinker identifies subject and object entities separately, whereas TPLinker mixes subjects and objects together and identifies them jointly. I also experimented with mixed identification in GPLinker, and found that the final performance showed no clear difference from separate identification.
3. For $S(s_h,o_h|p)$ and $S(s_t,o_t|p)$, TPLinker converts this into $l(l+1)/2$ separate 3-way classification problems, which introduces a pronounced class-imbalance issue; GPLinker, on the other hand, uses the multi-label cross-entropy I proposed, which avoids this imbalance problem and is easier to train. In fact, TPLinker's authors later became aware of this issue as well, and proposed TPLinker-plus, which also adopts this multi-label cross-entropy.
Of course, in my view, the main contribution of this post is not really these modifications that constitute GPLinker, but rather a "top-down" way of understanding joint relation extraction models: starting from the quintuple scoring function $S(s_h,s_t,p,o_h,o_t)$, analyzing its difficulties, and then simplifying it into the decomposition $\eqref{eq:factor}$ to "tackle it piece by piece." I hope this top-down reasoning process can offer readers some inspiration when designing models for more complex tasks.
Summary
This post shared a GlobalPointer-based joint entity-relation extraction model, "GPLinker," along with a "top-down" derivation for readers' reference.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.