Enhancing Document-level Relation Extraction by Entity Knowledge Injection^††thanks: National Natural Science Foundation of China (No. 61872172).

This module leverages coreference triples to exchange the contextual information between aliases, and thus the representations of alias mentions can be closer.

Coreference Distillation. A simple method is to model the coreference triples as a new type of edges in the document graph and reuse graph-based models [11, 13, 24]. However, such a method cannot be generalized to the sequence-based models since they do not construct document graphs. Furthermore, the accuracy of existing coreference resolution tools is still far from perfect, even they are trained on large corpora. To alleviate error accumulation, it is inappropriate to directly add the edges as strong guidance information in the RE models.

Knowledge distillation [6, 19], as a model compression technique and a solution to integrate external knowledge into a target model, has been used in a wide range of NLP tasks. In this paper, we leverage the idea of knowledge distillation and propose coreference distillation to inject coreference triples into the RE models. Our main idea is to leverage a pre-trained coreference resolution model which has been trained on a large coreference dataset as the teacher model, and then force the student model (i.e., the RE model) to generate a prediction probability distribution that approximates the teacher model on the coreference triples. Finally, the student model learns the coreference knowledge and generalization ability in the teacher model. Formally, for a coreference triple $(m_{s},m_{t},$ $p_{cr})$, its coreference probability generated by the teacher model is defined as

The student model generates the coreference probability with a multi-layer perceptron (MLP):

where $\mathbf{m}_{s}$ and $\mathbf{m}_{t}$ denote the hidden representations of alias mentions $m_{s}$ and $m_{t}$, respectively, which are calculated by averaging the hidden representations of tokens in $m_{s}$ and $m_{t}$, that is, $\mathbf{m}_{s}=\mathrm{avg}_{w_{j}\in m_{s}}(\mathbf{h}_{w_{j}}),\mathbf{m}_{t}=\mathrm{avg}_{w_{j}\in m_{t}}(\mathbf{h}_{w_{j}})$. “;” is the concatenation operation, and $\psi(m_{s},m_{t})$ denotes the shortest distance between $m_{s},m_{t}$ in the document. We divide the distance into $\{1,2,\dots,2^{\beta}\}$ bins, and associate each bin with a trainable distance vector. $\mathbf{\Delta}(\cdot)$ associates each $\psi$ to the distance vector of relevant bin. Empirically, aliases with different distances should have different impacts on each other. Therefore, we propose trainable distance vectors to model and utilize such difference.

We enforce the student model to learn from the teacher model using the following coreference loss:

where $\mathrm{KL}(\cdot)$ is the Kullback-Leibler divergence.

Context exchanging. Based on the learned coreference knowledge, we further enable each alias mention to interact with its most similar counterpart, so as to exchange the semantic information between them. Specifically, given an alias mention $m_{s}$, we update its hidden representation through $\mathbf{m}_{s}=\mathbf{m}_{s}+\mathbf{m}_{t^{*}}$, where $t^{*}=$ $\operatorname*{arg\,max}_{t}\big{\{}P_{\text{stu}}(m_{s},m_{t})\,|\,(m_{s},m_{t})\in C\big{\}}$. In this way, the representations of the pronouns in particular can be enriched via their referents.

Finally, we obtain the token representations enhanced by coreference knowledge through the representations of alias mentions (if exists):

In coreference encoding, the MLP contains $d_{\text{MLP}}(2d_{\text{token}}+d_{\text{dist}})+2d_{\text{token}}$ parameters, where $d_{\text{MLP}},d_{\text{token}},d_{\text{dist}}$ are the dimensions of MLP hidden layers, token representations and trainable distance vectors, respectively.

This module aims to encode the attribute triples and the relation triples to generate the KG representations of entities.

Attribute triple encoding. A KG defines a set of common attributes, e.g., rdfs:label and schema:description, to describe its entities. We encode the attribute triples in the KG and generate the attribute representations for corresponding entities. For each attribute triple of an entity, we concatenate the attribute name $a$ and attribute value $v$ into a token sequence $q=[a;v]=(w_{1},\dots,$ $w_{M})$. In order to cope with the out-of-vocabulary problem, we define a lookup function to convert each token to a token embedding:

where $\mathrm{WordEmb(\cdot)}$ returns the word embedding in GloVe, and $\mathrm{CharEmb(\cdot)}$ offers the average of character embeddings pre-trained with Skip-gram. Our method can work with other word or character embeddings easily.

Next, we leverage AutoEncoder to encode a sequence of token embeddings into an attribute triple embedding in an unsupervised way:

where AutoEncoder is pre-trained on the attribute triples. We conduct self-supervised training, and both encoder and decoder of AutoEncoder use BiLSTM. AutoEncoder has good capacity for feature extraction and compression. In our model, the input of AutoEncoder is a concatenation vector of an entity and its attributes. The reconstruction loss of AutoEncoder can help extract a better compressed feature representation while preserving the attribute knowledge.

Finally, we stack all attribute triple embeddings of an entity into a one-dimensional CNN to obtain the attribute representation of the entity:

where $\parallel$ denotes the stack operation, and $\mathbf{h}_{e_{i}}^{0}$ is the attribute representation of entity $e_{i}$, which would be used as the input representation for relation triple encoding below. Here, we choose CNN since the convolutional layer is a good feature extractor to learn high-level representations from value embeddings while reducing the dimension of output representations. Furthermore, we use the 1D convolution kernel as its invariance to the order of attribute embeddings.

Relation triple encoding. The relation triples present in the form of an entity-relation graph structure, and the topology and relation types are the key to encode such knowledge. Based on the attribute representations of entities, we employ a R-GAT [3] with $K$ layers to convolute the entity-relation graph. R-GAT incorporates relation types using different embeddings and calculates attention scores on all adjacent nodes based on entity embeddings and relation embeddings. Specifically, the node forward-pass update for the $(k+1)^{\text{th}}$ layer is

where $\mathbf{W}_{\text{in}}^{(k,b)}$ and $\mathbf{W}_{\text{out}}^{(k,b)}$ denote two trainable parameters of the $b^{\text{th}}$ attention head ($1\leq b\leq B$) at the $k^{\text{th}}$ layer. $\mathbf{h}_{i}^{(k)}$ and $\mathbf{h}_{j}^{(k)}$ are the node representations of entities $e_{i}$ and $e_{j}$ at the $k^{\text{th}}$ layer, respectively. $\mathbf{M}$ is a trainable mapping matrix corresponding to the relation types in the KG. $r_{ij}$ is the relation type between $e_{i}$ and $e_{j}$. $\mathrm{LeakyReLU}(\cdot)$ and $\sigma(\cdot)$ are the activation functions. $U_{i}$ is the neighbor set of $e_{i}$. In this way, the entity representations are updated via their all adjacent entity embeddings and relation embeddings at the previous layer.

We refer to the representations of entities after graph convolution as the KG representations of entities, which encode the knowledge in both attribute triples and relation triples. We simply denote the KG representation of entity $e_{i}$ by $\mathbf{h}_{e_{i}}$.

In the KG encoding, the attribute encoding has $d_{\text{Auto}}N_{\text{max}}N_{\text{kernel}}(d_{\text{kernel}}^{2}+1)$ parameters, where $d_{\text{Auto}}$ is the output dimension of AutoEncoder, $N_{\text{max}}$ is the maximum number of attributes that an entity has, $N_{\text{kernel}}$ is the number of kernels, and $d_{\text{kernel}}$ is the kernel size of CNN. The R-GAT network in the relation triple encoding contains $2(N_{\text{layer}}-1)N_{\text{head}}d_{\text{RGAT}}^{2}+d_{\text{RGAT}}d_{\text{ent}}$ parameters, where $N_{\text{layer}}$ is the number of layers in R-GAT, $N_{\text{head}}$ is the number of attention heads at each layer of R-GAT, $d_{\text{RGAT}}$ is the hidden dimension of R-GAT hidden layers, and $d_{\text{ent}}$ is the dimension of entity representations.

The token representations and the KG representations of entities capture different knowledge in independent semantic spaces. Following ERNIE [37] and K-BERT [10], we employ a representation reconciliation module to exchange the knowledge from entities in the KG with their linked tokens in the document.

The representation reconciliation module consists of $N$-stacked aggregators. For the $n^{\text{th}}$ ($1\leq n\leq N$) aggregator, given the token representations $\{\mathbf{h}_{w_{1}}^{n-1},\dots,$ $\mathbf{h}_{w_{J}}^{n-1}\}$ and the KG representations of entities $\{\mathbf{h}_{e_{1}}^{n-1},\dots,\mathbf{h}_{e_{I}}^{n-1}\}$ from the preceding aggregator, the fusion phase is formulated as

where $\tilde{\mathbf{h}}_{w_{j}}^{(n)},\tilde{\mathbf{h}}_{e_{i}}^{(n)}$ are the token representation and KG representation after the multi-head self-attention [22], respectively. $\tilde{\mathbf{W}}_{w}^{(n)},\tilde{\mathbf{W}}_{e}^{(n)},\tilde{\mathbf{b}}^{(n)}$ are three trainable parameters. The information in the two semantic spaces is mutually integrated.

Then, the reconstruction phase leverages $\tilde{\mathbf{h}}$ to refine the output representations of each token and entity in the aligned token-entity pairs:

Here, the aligned token representations and entity representations are updated and enhanced by the integrated information. Note that the representations of entities without aligned tokens would not be updated.

Finally, we obtain the token representation sequence $\{\mathbf{h}_{w_{1}}^{N},\dots,\mathbf{h}_{w_{J}}^{N}\}$ from the last aggregator, which would constitute $\mathbf{H}^{\prime}$ and be fed to the prediction layer.

To supervise the above process, we employ a token-entity alignment task. For each aligned token-entity pair $(w_{j},e_{i})$, we predict the aligned KG entity $e_{i}$ based on the token $w_{j}$. We only ask the model to predict entities within a given entity candidate set. By default, all linked entities in the document form the candidate set. For the token sequence $\{w_{1},$ $\dots,w_{J}\}$ and the corresponding candidate entities $\{e_{1},\dots,$ $e_{I}\}$, the token-entity alignment loss is

where $f_{j,i}^{*}\in\{0,1\}$ is the true alignment label between $w_{j}$ and $e_{i}$, and $P(e_{i}\,|\,w_{j})=\frac{\mathrm{exp}\big{(}\textrm{Linear}(\mathbf{h}_{w_{j}}^{(N)})\cdot\mathbf{h}_{e_{i}}\big{)}}{\sum_{l=1}^{I}\mathrm{exp}\big{(}\mathrm{Linear}(\mathbf{h}_{w_{j}}^{(N)})\cdot\mathbf{h}_{e_{l}}\big{)}}$ returns the probability that $e_{i}$ can be predicted by $w_{j}$.

We optimize the RE loss, coreference loss and token-entity alignment loss with multi-task learning. The final loss is

where $\alpha_{1},\alpha_{2}$ and $\alpha_{3}$ are the weight hyperparameters.

In the representation reconciliation, for each aggregator, the multi-head self-attention networks contain $4d_{\text{token}}^{2}+4d_{\text{ent}}^{2}$ parameters, the fusion phase contains $d_{\text{out}}(N_{\text{token}}+N_{\text{align}})+N_{\text{token}}$ parameters, and the reconstruction phase contains $2N_{\text{align}}(d_{\text{out}}+1)$ parameters, where $d_{\text{out}}$ is the output dimension of multi-head self-attention networks, $N_{\text{token}}$ is the number of tokens, and $N_{\text{align}}$ is the number of aligned token-entity pairs. Therefore, the parameters of representation reconciliation are $N_{\text{agg}}\big{[}4d_{\text{token}}^{2}+4d_{\text{ent}}^{2}+d_{\text{out}}(N_{\text{token}}+N_{\text{align}})+N_{\text{token}}+2N_{\text{align}}(d_{\text{out}}+1)\big{]}$, where $N_{\text{agg}}$ is the number of aggregators.

Model complexity. The total parameter number of KIRE is $d_{\text{MLP}}(2d_{\text{token}}+d_{\text{dist}})+2d_{\text{token}}+d_{\text{Auto}}N_{\text{max}}(d_{\text{kernel}}^{2}+1)+2(N_{\text{layer}}-1)N_{\text{head}}d_{\text{RGAT}}^{2}+d_{\text{RGAT}}d_{\text{ent}}+N_{\text{agg}}\big{[}4d_{\text{token}}^{2}+4d_{\text{ent}}^{2}+d_{\text{out}}(N_{\text{token}}+N_{\text{align}})+N_{\text{token}}+2N_{\text{align}}(d_{\text{out}}+1)\big{]}.$

Datasets. We select two benchmark datasets in our experiments: (1) DocRED [32] is a crowdsourced dataset for document-level RE. The relation labels in its test set are not public. (2) DWIE [33] is a new dataset for document-level multi-task information extraction. We use the data relevant to RE only. Since DWIE does not have the validation set, we randomly split its training set into 80% for training and 20% for validation. Table 1 lists the statistical data.

Knowledge graph. We select Wikidata (2020-12-01) as our KG due to its coverage and popularity [2, 21]. The numbers of its relation and attribute triples are 506,809,195 and 729,798,070, respectively. To prevent the test leakage, we filter out all the relation triples with entity pairs to be labeled in the test sets.

Evaluation metrics. We measure F1-score (F1) and Ignore F1-score (Ign F1) in our experiments. We repeat five times using the same hyperparameters but different random seeds, and report the means and standard deviations.

Implementation details. To achieve good generalization, we do not carry out excessive feature engineering on Wikidata. Numerical attributes are regarded as texts, and their semantics are captured by the word embeddings [14]. We employ NeuralCoref 4.0 as our coreference resolution tool and also use the annotations provided in DocRED and DWIE. We use a two-stage method for training, which first trains a basic RE model and then fine-tunes this model to train the knowledge injection layer. The training procedure is optimized with Adam. Moreover, to compare fairly, the basic RE model and its corresponding KIRE adopt the same hyperparameter values. We set the batch size to 4 and the learning rate to 0.0005. We use three R-GAT layers and two aggregators. Moreover, $\alpha_{1},\alpha_{2},\alpha_{3}$ are 1, 0.01 and 0.01, respectively. The dimension of hidden layers in MLP is 256, the dimensions of GloVe and Skip-gram are 100, and the dimension of hidden layers in AutoEncoder is 50. See the source code for more details.

Improvement on baseline models. To validate the effectiveness and versatility of KIRE, we pick four baseline models in [32]. The first three models directly employ CNN, LSTM and BiLSTM to encode documents, while the fourth model is called context-aware, which leverages the attention mechanism with BiLSTM. These four models are native to the DocRED dataset and widely chosen as the competitors in many RE studies [7, 11, 18, 27, 28, 30, 35, 38].

Table 2 depicts the result improvement, and we observe that: (1) KIRE consistently improves the performance of all baselines on DocRED and DWIE, which demonstrates the good generalization of KIRE. Small standard deviations also tells the good stability of KIRE. (2) KIRE obtains a significant improvement of Ign F1/F1 up to 2.63/2.47 on DocRED and 2.03/1.95 on DWIE, respectively. This is mainly because the ways that the baseline models encode a document are too simple to capture some part of important contextual information in the document. External knowledge from KIRE makes up for this part, and therefore effectively improves the model performance. (3) CNN performs poorly, because the text order is important for RE while CNN cannot process such order well.

Comparison with existing knowledge injection models. We choose three recent models: RESIDE [21], RECON [2] and KB-graph [25], which inject extra knowledge into RE models. Specifically, we use KB-graph instead of the full version KB-both since it selects Wikipedia as another knowledge source, which is unfair to other models. To compare fairly, we only adopt the knowledge injection modules of the above models to enhance the token representations in the documents, and the representations are used by the baseline RE models to predict the relation labels.

Table 2 presents the comparison results, and we obtain several findings: (1) KIRE is consistently superior to RESIDE, RECON and KB-graph with an improvement of Ign F1/F1 up to 0.71/0.66 on DocRED and 0.39/0.43 on DWIE, respectively. Given that the test sets contain (ten) thousand relation instances, we think that the improvement makes sense. For example, on the validation set of DocRED, KIRE can correctly predict an average of 478 more instances than the second best method RECON. Such improvement brought by KIRE attributes to that KIRE absorbs more knowledge like coreferences and fuses the knowledge better. (2) The improvement brought by RESIDE is the lowest since it only injects limited knowledge like entity types and relation aliases. RECON and KB-graph explore more knowledge from the KG, but they still ignore the coreference knowledge. Besides, the methods that they employ to integrate knowledge are representation average or concatenation, which may lose part of semantic information in the injected knowledge.

Improvement on state-of-the-art models. We employ two sequence-based models, ATLOP [38] and SSAN [30], as well as a graph-based model, GLRE [27], due to their good performance and open source. Enhancing these models is very challenging, since they have already explored various information in the documents and achieved state-of-the-art results. Due to the limit of GPU RAM, we use the BERT-base versions of ATLOP, SSAN and GLRE and re-run them according to the hyperparameters reported in their papers and source code.

The result improvement is shown in Table 3, and we have several findings: (1) For the two sequence-based models, KIRE obtains an improvement of Ign F1/F1 up to 1.03/0.95 on DocRED and 3.08/1.81 on DWIE, respectively. This mainly attributes to the fact that the extra knowledge injected by KIRE can effectively help the models identify and capture more interactions between entity pairs especially across sentences. (2) For the graph-based model, KIRE obtains an improvement of Ign F1/F1 up to 1.14/1.09 on DocRED and 1.37/1.11, respectively. This is largely due to the fact that the extra knowledge injected by KIRE can enrich the representations of mention nodes and entity nodes in the document graphs for more accurate reasoning between entity pairs especially of longer distance. (3) This also verifies that our knowledge injection framework can be generalized to a broad range of document-level RE models.

Ablation study. We conduct an ablation study on the four baseline models. For “w/o distill”, we disable the coreference distillation module and directly use the original coreferences as the injected knowledge. For “w/o attr.”, we initialize the relation triple representations by max pooling the word embeddings of entity labels. For “w/o rel.”, we directly adopt the attribute representations of entities as KG representations. For “w/o KG”, we disable the whole KG encoding module. Additionally, we replace KIRE with three simple variants for knowledge injection. For “w/ rep. avg”, we average the hidden representations of alias mentions, and the token representations are averaged with KG representations of entities. For “w/ rep. concat”, we concatenate the representations of alias mentions, and the KG representations of entities are concatenated after the aligned token representations. For “w/ MLP”, we leverage two MLP layers to fuse the representations of alias mentions and the KG representations of entities with the aligned token representations, respectively.

From Figure 3, we can see that: (1) Ign F1/F1-scores reduce when we disable any modules, showing their contributions. (2) The changes caused by removing one type of knowledge are not obvious, mainly due to the crossovers among the three types of knowledge in the information space. (3) The results decline if we disable the coreference distillation, due to the coreference errors in the injected knowledge. (4) If we remove the KG encoding, the results drop drastically, as the baseline models cannot generate extra relation and attribute knowledge. (5) Compared to the three variants, the larger increase brought by KIRE validates the effectiveness of coreference distillation and representation reconciliation.

Influence of mention number. We measure the effectiveness of KIRE w.r.t. average mention number for each entity pair. For DocRED, we evaluate it on the validation set. The results are shown in Table 4. We observe that KIRE gains higher performance for the entity pairs with more mentions, in particular when the average mention number $>3$. This is because KIRE injects knowledge into the RE models by updating the token representations of entity mentions, which has a greater impact on the entities with more mentions.

Comparison with alternative graph encoders. We compare R-GAT with GCN [9], GAT [23] and R-GCN [17]. We remove the coreference encoding and attribute triple encoding to eliminate their interference. From Figure 4, the performance of R-GCN and R-GAT is better than GCN and GAT, as they can capture the relation information in the entity-relation graphs. The results of GAT are greater than GCN, as GAT can selectively aggregate the neighboring information by self-attention. Similarly, R-GAT slightly outperforms R-GCN.

Case study. We depict two successful cases and a failed case in Table 5. We still use CNN, LSTM, BiLSTM and Context-aware as baselines.

• •

  Case 1. To identify the relation between North Dumfries in $S2$ and Regional Municipality of Waterloo in $S1$, we use the extra knowledge (Regional Municipality of Waterloo, instance of, regional municipality of Ontario) and the coreference of It and The Waterloo Moraine from NeuralCoref to correctly infer the relation P131 (located in the administrative territorial entity).

• •

  Case 2. With the aid of the extra knowledge (Meghan Markle, country of citizenship, United States of America) and the coreference of Meghan Markle and Markle from NeuralCoref, we successfully detect the relation citizen_of between Thomas Markle in $S10$ and America in $S1$.

• •

  Case 3. To recognize the relation between Ensign George Gay in $S2$ and the United States Navy in $S4$, we require a bridge entity Robert Kingsbury Huntington. Through the coreference of Robert and He from NeuralCoref, we identify that Robert and George are comrades. Then from $S4$, we find out that Robert is enlisted in the US Navy. According to this reasoning chain, we can see that the relation between George and the US Navy is P241 (military branch). KIRE fails to run such complex reasoning involving three sentences.