KnowRA: Knowledge Retrieval Augmented Method for Document-level Relation Extraction with Comprehensive Reasoning Abilities

1 Introduction

The document-level relation extraction (Doc-RE) task aims to extract pre-defined relation triples from documents containing multiple sentences. Compared with the existing sentencelevel RE task, Doc-RE is not only more complicated but also more fundamental for real-world applications, like retrieval augmented generation (RAG) method for large language model (LLM) [Dai et al., 2022; Yang et al., 2023;

Xu et al., 2024], automatic question and answering [Ding et al., 2022; Orogat and El-Roby, 2022], and event extraction [Man et al., 2022; Huang et al., 2023; Yuan et al., 2023].

Obviously, Doc-RE models need to have comprehensive reasoning abilities for Doc-RE. Previous studies divided these abilities into four categories [Yao et al., 2019]: pattern recognition, co-reference reasoning, common-sense reasoning, and logical reasoning, as shown in Figure 1.

However, the existing Doc-RE models were mostly optimized for partial reasoning abilities but lacked comprehensive reasoning abilities. For pattern recognition, most recent methods constructed graph structures [Xu et al., 2021c; Wei and Li, 2022; Peng et al., 2022; Zhang et al., 2023b] to establish long-distance associations between entities. For logical reasoning, several models used bridge/intermediary entities or evidence sentences to establish relations between two entities that are not directly connected [Zhang et al., 2023a; Huang et al., 2024; Xu et al., 2022]. Considering that entity mentions often appear in the form of co-reference pronouns in the document, methods based on co-reference reasoning have been proposed [Ye et al., 2020] and focused on identifying co-reference pronouns of entities, which were used as bridge entities to establish indirect relations between entities.

Also, the Doc-RE model requires the ability of commonsense reasoning because some relations cannot be directly inferred from the document itself, and external knowledge is needed to assist in Doc-RE. However, few existing Doc-RE models combined external knowledge for common sense reasoning [Wang et al., 2022a].

To summarize, the main challenges for Doc-RE lie in: 1) How to integrate the comprehensive reasoning abilities required by Doc-RE. 2) How to represent and integrate external knowledge with the internal semantics of the document. 3) How to autonomously determine whether to accept external knowledge, considering that external knowledge may be lagging, one-sided, or even wrong.

To solve above mentioned problems, we proposed a comprehensive reasoning method for Doc-RE with knowledge retrieval augmentation and filtration, as shown in Figure 2. Firstly, a heterogeneous multi-level document graph was constructed for semantic encoding of entities, mentions, and sentences of the document. Then, a pre-trained co-reference resolution model was introduced to establish associations between entities and their corresponding pronouns for coreference reasoning. Moreover, the document graph was extended with the retrieved external knowledge for commonsense reasoning. On this basis, a novel knowledge filtration method was proposed to determine whether to accept external knowledge. Finally, an axial attention method was proposed to use intermediary entities for realizing logical reasoning across multiple sentences.

The main contributions of our work are three folds:

• A comprehensive method for Doc-RE was proposed by achieving document semantic encoding with a multilayer heterogeneous document graph, integrating the coreference resolution model, injecting external knowledge, and introducing the axial attention mechanism.
• A knowledge augmentation and filtration method was proposed to judge whether to accept external knowledge, for filtering out irrelevant knowledge.
• Extensive experiments performed on two public datasets demonstrated the superiority of our method compared to the state-of-the-art (SOTA) baselines.

2 Our Methodology

2.1 Construction of document graph

Firstly, a pre-trained language model was employed to perform semantic encoding operations on the input document. Then, a multi-level heterogeneous document graph was present to model the connections among different entities, mentions, sentences in a document, defined as follows:

Definition 1. Multi-level Heterogeneous Document Graph (MHDG). MHDG=< V,E >, where V = {v|v ∈ V M ⋃ V S ⋃ V D } represents the node set and E = {< vi, vj > |vi, vj ∈ V, i ̸= j}, which represents edges between nodes vi and vj .

According to definition 1, we defined three types of node in MHDG: Mention node (V M ), Sentence node (V S ), and Document node (V D ). Then, four types of edges were given:

1. Document-Sentence Edge: EDS = {< vD i , vS j > |vD i ∈ V D , vS j ∈ V S }. The document node and all sentence nodes are connected through these edges.
2. Sentence-Sentence Edge: ESS = {< vS i , vS j > |vS i , vS j ∈ V S , i ̸= j}. Two adjacent sentences are connected by sentence-sentence edges.
3. Mention-Sentence Edge: EMS = {< vM i , vS j > |vM i ∈ V M , vS j ∈ V S }, connecting mention nodes and the sentence node appearing with the same one sentence.
4. Mention-Mention Edge: Two mention nodes are connected by the Mention-Mention Edge (MME), which can be further divided into two categories: Co-Occurrence Mention-Mention Edge (CO-MME) means different mentions appear in the same sentence, formalized as: ECO-MME = {< vM i , vM j > |vM i , vM j ∈ V M }. CoReference Mention-Mention Edge (CR-MME) means that different mentions refer to the same entity, formalized as: ECR-MME = {< vM i , vM j > |vM i , vM j ∈ V M }.

Then, the semantic representation was performed on the MHDG with a graph-based neural network, denoted as:

where Tinput represents the input tokenized sequence, L represents the length of the document. The representations of three kinds of nodes are as follows:

where HD ∈ Rd represents the d dimension semantic representation of the document node. Hmij represents the semantic embedding of mention node for the j-th mention of the i-th entity, and P i,j e represents the position of the j-th mention of the i-th entity, and "*" denotes for the special symbol placed before the mention mij . HSi represents the semantic representation of sentence node Si.

Then, the Graph Attention Network (GAT) was used to calculate association score AS(ei,j ) for each edge between node i and j:

[S1] Dwight Tillery is an American politician of the Democratic Party… [S2] Tillery grew up in the city's West End … [S3] … [S4] Tillery served as mayor of Cincinnati from 1991 to 1993. [S5] … [S6] … [S7] ... supporting the election campaigns of African American candidates.

Figure 2: Overview of the proposed comprehensive reasoning method with knowledge augmentation and filtration for Doc-RE.

where ⊕ represents the concatenation operation. The updated node representations based on MHDG are as follows:

where Wα ∈ Rd and Wβ ∈ Rd×d are trainable parameters. N(i) represents the number of adjacent nodes to node i.

2.2 Co-reference reasoning component for Doc-RE

To solve the co-reference reasoning problem, we firstly pretrained a co-reference resolution model, Mcoref , and then used it to identify the co-reference pronouns, as follows:

where Ci represents the set of identified co-reference pronouns and nci represents the number of the co-reference pronouns. Then, we used attention matrix between each input token and recognized co-reference pronouns to update semantic representations of co-reference pronouns, as follows:

where A ∈ RH×L×L is the attention matrix of all input tokens, H represents the number of attention heads, and L represents the length of input tokens. Qi ∈ RL, represents the averaged vector of attention value of entity Ei to its coreference pronouns. Ah,k i represents the h-th attention head of entity Ei to its k-th co-reference word. P i,k c represents the position of each co-reference pronoun.

Finally, we obtained the semantic representation of cik as shown in formula (11), where Q′i [P i,k c ] ∈ Rnci,k represents the normalized attention value of entity Ei to its k-th coreference pronoun. hP i,k c represents the semantic embedding of the co-reference pronoun ci,k.

2.3 Common-sense reasoning with knowledge augmentation and filtration

Construction of knowledge-augmented document graph

To introduce external knowledge for common-sense reasoning, we further constructed a knowledge-augmented document graph, denoted as K-MHDG.

Definition 2. Knowledge-retrieval-augmented Multi - level heterogeneous document graph (K-MHDG). KMHDG=< V, E ∪Eknow >, where V = {v|v ∈ V {M,S,D} } represents the entity node set, E ⊆ {< vi, vj > |vi,j ∈ V, i ̸= j} represents the relation edges between node vi and vj in MHDG, and Eknow represents the newly added edges, consisting of relations between entities, denoted as rknow, retrieved from the external knowledge base, formalized as: Eknow ⊆ {< vm, vn > |∃erknow =< vm, vn >∈ hasRel(QEi id , rknow,QEj id ), vm,n ∈ V M } .

We selected Wikidata as the external knowledge base to construct the K-MHDG. By utilizing the interface getQid(EntityName) provided by Wikidata, all possible entity identifiers, i.e., Qids, were obtained based on the corresponding entity name. Then, based on the hasRel(Qtid , rknow,Qhid ) interface in Wikidata, we can query all relation triples with the provided entity pairs. Finally, the newly retrieved relation ernew =< vM Enti , vM Entj > would be added to Eknow and the MHDG was augmented to K-MHDG.

Semantic encoding for the external knowledge

Based on K-MHDG, the representation of entity node Ei was updated as follows:

where ⊕ represents the concatenation operation, hmij represents the embedding of mention mij and nmi represents the number of mentions related to Ei. hcik and nci refer to the semantic embedding and number of the co-reference pronouns for entity Ei, respectively.

External knowledge filtration method

Considering that the external knowledge introduced is not always correct, we further proposed a new confidence-scorebased knowledge filtration method to help our model autonomously determine whether to accept external knowledge.

The confidence score for each edge, i.e., ei,j =< Ei, rk, Ej >, was denoted as τi,j and calculated as:

where diag(·) function represents the diagonal matrix with the vector rk as the diagonal element. The confidence score represents the likelihood of a relation existing between entity Ei and Ej. After calculating the confidence score, the entity representation of entity Ei was also updated as follows:

where σ(·) represents the sigmoid function, which is used to convert confidence score into probabilities with values between 0 and 1. ri,j represents the set of relations that exist between entities Ei and Ej .

Then, we used the correct relation label set provided by the annotated training set to train the model in a way that increases the confidence score for correct relations and decreases the confidence score for incorrect relations. The optimization function is:

(15)

where N represents the edge number in K-MHDG, and yn represents the annotated relation label in the training set, which value is 1 (correct relation) or 0 (incorrect relation).

2.4 Logical reasoning with axial attention

Semantic fusion based on common context of entity pair

Because Doc-RE is implemented based on the entity pairs, it is necessary to integrate context information into entity pairs for logical reasoning, not just context of individual entities.

Firstly, we used the attention mechanism proposed to obtain the context of the common concern of entity Ei and Ej, that is, the relevant context representation of entity pair < Ei, Ej >. The formula was given as follows:

where symbol "×" represents the outer product of embeddings, and "·" represents the dot product of embeddings.

Then, the semantic embedding for head and tail entities fused with the context of entity pair < Ei, Ej > was calculated as:

where zi and zj represent the entity embedding fused with the context information. W∗ ∈ Rd represents the trainable model parameters. Finally, the semantic representation of entity pair < Ei, Ej >, denoted as zi,j , was obtained.

Axial-attention-based method for logical reasoning across sentences

We arranged all entity pair representations zi,j in the document into an entity pair matrix of N × N, where N = |V | represents the number of entities. If there is an intermediary entity Ek between entity pair < Ei, Ej >, the purpose of our proposed method is to establish a multi-hop logical reasoning model by integrating semantic representations of all intermediary entities and the corresponding entity pairs.

In the entity pair matrix, the semantic information of the horizontal intermediary entity pair < Ei,Ek|k = 1, . . . ,N > in the same row was firstly fused into < Ei,Ej >, calculated as:

where qi,j, ki,j, vi,j = Wqzi,j,Wkzi,j,Wvzi,j is the query, key, and value vector obtained by linear transformation of zi,j. Wq,Wk,Wv ∈ Rd×d are query, key, and value vectors obtained by linear transformation.

Then, the similar operation was performed on all vertical entity pairs in the same column and calculated as follows:

Finally, all intermediary entity pairs < Ei,Ek > and < Ek,Ej > were fused and the final representation for < Ei, Ej > was obtained by :

Adaptive Relation Extraction Loss for Doc-RE

Considering Doc-RE is a multi-label classification task, we adopted adaptive threshold loss function [Zhou et al., 2021; Guo et al., 2023] with a threshold class, denoted as TH, to adaptively separate positive (R+) and negative (R-) relations:

where Li,j represents the probability that entity pair < Ei, Ej > belongs to each predefined relation type:

Through the joint training of the loss function LRE and LKRA, the final optimization function for Doc-RE is:

where λ is the pre-defined hyper-parameter.

3 Experiments and Analyses

3.1 Experimental Settings

Datasets. We evaluated our model on two public datasets for document-level RE. Re-DocRED [Tan et al., 2022b] is a high-quality revised version of DocRED [Yao et al., 2019]. Re-DocRED corrects the false negatives problem in dataset DocRED and contains 3,053 documents for training, 500 for development, and 500 for the test set. DWIE [Zaporojets et al., 2021] is sampled and annotated from the news website Deutsche Welle, containing 602, 98, 99 documents for training, development, and testing, respectively, with 43,373 entities, 21,749 relational facts, and 65 relation types.

Evaluation metrics. We used the micro F1 (F1), ignore F1 (Ign F1), Intra F1, and Inter F1 as the metrics for model performance, following previous work [Wang et al., 2022b]. Ign F1 is a revised version of F1, which excludes the shared relations between the training and development/test set. Intra F1 is used to evaluate F1 of relation triples that appear in the same sentence. Inter F1 is used to evaluate F1 of crosssentence relation triples.

3.2 Baselines

According to different model structures, the following models were selected for performance comparison.

• Sequence-based models: These models used different deep neural structures for relation extraction, including convolution neural network (CNN), Bi-LSTM, and Context-Aware LSTM [Yang et al., 2023].
• Graph-based models: These models used graph neural network for Doc-RE, including GAIN [Zeng et al., 2020], SIRE [Zeng et al., 2021], DRN [Xu et al., 2021b], SagDRE [Wei and Li, 2022].
• Transformer-based models: These Transformer-based models [Vaswani et al., 2017] include SSAN [Xu et al., 2021a], ATLOP [Zhou et al., 2021], ATLOP-MILR [Fan et al., 2022], DocuNET [Zhang et al., 2021a], KDDocRE [Tan et al., 2022a], UGDRE [Sun et al., 2023], and JMRL-DREEAM [Qi et al., 2024]. It should be

noted that documents in dataset DWIE are quite long, with over 50% of documents exceeding 768 tokens, and the maximum length reaches 2560 tokens. Therefore, for dataset DWIE, we replaced RoBERTa-large with LongFormer [Beltagy et al., 2020] as the backbone model, which supports a maximum of 4096 tokens. For dataset Re-DocRED, baseline models used RoBERTa-large as the cornerstone model.

• LLM-based models: We used 13B LLAMA-2 1 model as the large language model for Doc-RE, finetuned with LoRA 2 .

Considering that these baseline models used different datasets and cornerstone models in their original papers, we tried our best to rerun and fine-tune these models for fair comparison. For example, the rerun model JMRL-DREEAM even achieved better performance on dataset Re-DocRED compared to its original paper.

3.3 Main Results

Table 1 listed the performance of models on two datasets. We observed that: 1) Our KnowRA model outperformed other baseline models in almost metrics on two datasets. For dataset Re-DocRED and DWIE, KnowRA surpassed the existing SOTA model, JMRL-DREEAM, by 0.28 and 1.13 in F1 score, respectively. These experiments proved the superiority of our model for Doc-RE. 2) In terms of Intraand Inter-F1, our model achieved the second-best results in dataset Re-DocRED and the best performance in dataset DWIE. These experimental results proved the advantages of our model in cross-sentence relation extraction. Meanwhile, for dataset DWIE with longer document length, our method achieved better performance compared to the SOTA models, i.e., KD-DocRE and JMRL-DREEAM, which indicates that our model has stronger semantic reasoning ability for extracting long-distance document-level relations. 3) For both datasets, our model outperformed the LLaMA-2-based model in all metrics, which implies that large language models need to improve their effectiveness through relevant optimization technologies when applied to downstream tasks.

3.4 Ablation Experiments

Ablation experiments were shown in Table 2. For dataset ReDocRED, after removing the document graph (denoted as "- w/o graph") and axial attention (denoted as "-w/o axial atten."), the model performance decreased by 1.07 and 1.01 in F1 score, respectively. When both the two components were removed (denoted as "-w/o graph+axial"), the performances on F1 decreased by 1.28. The removal of the co-reference reasoning component (denoted as "-w/o coref ") also leads to a drop of 0.70 in terms of F1.

In addition, the common-sense reasoning component based on the knowledge-retrieval-augmentation (denoted as "-w/o knowAug.", representing the removal of the K-MHDG) is the most important component, as the F1 score decreased the most (-1.10 on Re-DocRED and -1.02 on DWIE) when it was removed. Also, when the knowledge filtration method (denoted as "-w/o knowFilt.") was removed from our complete model, the F1 score decreased by 0.97 and 0.79 on dataset Re-DocRED and DWIE, respectively, which proves that the filtration of the external knowledge is necessary for filtering out noise information and boosting performance. For dataset DWIE, we observed similar phenomena, which verifies the effectiveness of our comprehensive reasoning components.

1 https://github.com/meta-llama/llama.

2 https://github.com/microsoft/LoRA.

3.5 Effects of Intra- and Inter-sentence Reasoning

The number of sentence intervals was used to present the distance between the head and tail entities in a relation triple. According to the results shown in Figure 3, we found that: 1) When the head and tail entities are in the same sentence (intervals = 0), our model performed better than the baselines. 2) When the head and tail entities are from different sentences (intervals > 0), our model was also superior to other baselines. 3) The performance of all models gradually decreased with the increase of sentence intervals, which indicates that our model performed well in both intra- and inter-senetence reasoning.

3.6 Effects of Context Length

We conducted experiments to quantitatively evaluate the impact of the context length on the performance of Doc-RE. For dataset DWIE, the length of most documents (94.81%) exceeds the maximum length supported by the RoBERTalargebased encoder. To this end, we replaced the original encoders with longFormer [Beltagy et al., 2020].

The experimental results are shown in Figure 4 and we found that: 1) Longer context can greatly improve the performance of Doc-RE. 2) As the context length gradually increases, the performance of all models first improves, and then the growth rate gradually slows down. Similar trends were also observed on other baselines. These results proved that our model can outperform SOTA models in different context length, which indicates a good scalability for our model.

3.7 Impact of Knowledge-augmentedMethod

We conducted a hyper-parameter analysis on weight λ of the LKRA loss function to quantitatively evaluate the impact of the proposed knowledge-retrieval-augmentation method.

As shown in Figure 5, it can be observed that: 1) For both datasets, the model performance shows a trend of first improving and then decreasing with the increase of parameter λ. 2) Introducing external knowledge through our proposed knowledge augmentation method can improve the model performance for Doc-RE. However, too high the weight value λ would lead to a decline for the model performance. One possible explanation is that external knowledge may be wrong or inconsistent with the internal information of the document.

3.8 Case Study

Figure 6 shows the construction of K-MHDG by introducing the common-sense knowledge from the external knowledge base. The nodes in the figure represent entities, where the numbers indicate the identification index of the entities. The directed edges between nodes represent the relations between entities, the words on the edges represent the relation types, and numbers on edges represent the confidence score, i.e., τ.

Our knowledge augmentation method can extend the relation triples that were not annotated in the original dataset, identify the correct relation triples, denoted by the green ✓, and also filter out noise external knowledge and incorrect relation triples, which are denoted by the red ×.

4 RelatedWork

For pattern recognition and logical reasoning, existing models mainly used intermediary entities or evidence sentences to achieve logical reasoning [Yuan et al., 2023; Zhang et al., 2021b; Wei and Li, 2022]. However, evidence sentences are not always available. Peng et al. [2022] proposed a subgraph seasoning model for Doc-RE, which places emphasis on key entities and integrates various paths between entity/mention pairs into a subgraph for relation reasoning.

For co-reference reasoning, this type of model has positive effects on improving the performance of modeling long-distance reasoning by reducing ambiguity between coreference entities and mentions involving multiple sentences [Xue et al., 2022], Ye et al. [2020] proposed mention reference prediction method to equip the language model with the capacity for capturing the co-reference relations. Wang et al. [Wang et al., 2022a] distilled the co-reference reasoning ability into the relation extraction model.

However, for common-sense reasoning, using only the information in the current document makes it difficult to establish implicit associations and determine relation types between different entities/mentions in multiple sentences, without the help of common sense distilled from external knowledge [Li et al., 2021]. Existing models rarely incorporate external knowledge into the Doc-RE model. To solve this problem, we introduced external knowledge into our model with the knowledge filtration method.

5 Conclusions

In this paper, we presented a comprehensive reasoning reasoning model for Doc-RE. Concretely, the proposed knowledge document graph and three different reasoning components, i.e., graph-based semantic encoding, co-reference resolution, knowledge augmentation and filtration, and axial attention method were integrated into our KnowRA model to enhance the comprehensive reasoning abilities. Experiments conducted on two benchmarks datasets demonstrated the superiority of the proposed model in Doc-RE.

Acknowledgments

This work was supported by the National Natural Science Foundation of China (Nos. 61572250 and 62476135), Jiangsu Province Science & Tech Research Program (BE2021729), Open project of State Key Laboratory for Novel Software Technology, Nanjing University (KFKT2024B53), Jiangsu Province Frontier Technology Research and Development Program (BF2024005), Nanjing Science and Technology Research Project (202304016) and Collaborative Innovation Center of Novel Software Technology and Industrialization, Jiangsu, China.