A Densely Connected Criss-Cross Attention Network for Document-level Relation Extraction

We first treat the document $D$ as a sequence of words, i.e. $D\!=\!\{w_{i}\}_{i}^{L}$, where $L$ is the total number of words in the document. Then, we insert special symbols $\left\langle E_{i}\right\rangle$ and $\left\langle\setminus E_{i}\right\rangle$ to mark the start and end positions of the mention respectively, where $E_{i}$ is the entity id of the mention. It is an improved version of entity marker technology Zhou et al. (2021); Zhang et al. (2021); Shi and Lin (2019); Soares et al. (2019) by introducing entity id information which can help align the information of different mentions from the same entity. Finally, we leverage the pre-trained language model as an encoder to convert documents $D$ into a sequence of contextual embeddings as follows:

We take the embedding of $\left\langle E_{i}\right\rangle$ at the start of mentions $m_{j}$ as the mention embedding $h_{m_{j}}$. Then, we leverage logsumexp pooling Jia et al. (2019), a smooth version of max pooling, to obtain the embedding $h_{e_{i}}$ of entity $e_{i}$ which contains $p_{i}$ mentions $\{m_{j}^{i}\}_{j}^{p_{i}}$:

After obtaining the embedding of all entities in the document, we construct an Entity-Pair Matrix $M{\in}R^{N_{e}\times N_{e}\times d}$ through the Relation Feature Calculation module, where $N_{e}$ refers to the number of entities and $d$ is the dimension of the relation feature vector. The $M_{s,o}$ item in $M$ represents the relation feature vector between the entity $e_{s}$ and the entity $e_{o}$, which is calculated as follows:

where $h_{e_{s}}$ is subject entity embedding, $h_{e_{o}}$ is object entity embedding, $h_{doc}$ is document embedding, and $c_{s,o}$ is entity-pair-aware context feature, $FFN()$ refers to a feed-forward neural network, $W_{o},W_{s}$ is the learnable weight matrix.

We use the embedding of the document start token “[CLS]” as the document embedding $h_{doc}$, which can help aggregate cross-sentence information and provide document-aware representation.

The entity-pair-aware context feature $c_{s,o}$ represents the contextual information in the document that the entity $e_{s}$ and the entity $e_{o}$ pay attention to together. The $c_{s,o}$ is formulated as follows:

where $A_{s,i}$ is the attention score of the entity $e_{s}$ paying attention to the i-th token $h_{i}$ in the document.

In this part, we introduce the Dense-CCNet module in detail. As shown in Figure 2, the Dense-CCNet module consists of densely connected $N$ identical layers that are composed of two sub-modules: the Criss-Cross Attention (CCA) module and the Transition module.

We followed the CCNet model Huang et al. (2019) for the CCA module. Each entity pair in the entity-pair matrix can pay attention to the relation feature of other entity pairs in horizontal and vertical directions through the CCA module. The CCA module can be formulated as follows:

where $A_{(s,o)\rightarrow(s,i)}$ is the attention score of the $M_{s,o}$ paying attention to the $M_{s,i}$. Therefore, the CCA module can complete entity-pair-level one-hop reasoning on the entity-pair matrix, and it is possible to complete multi-hop reasoning by stacking multiple layers of the CCA module.

However, simply using Recurrent Criss-Cross Attention (RCCA) Huang et al. (2019) to complete the logical reasoning of the entity-pair level may have several problems: (1) The RCCA only focuses on the high-level multi-hop inference feature and ignores the low-level single-hop inference feature which is also very important for document-level RE. (2) The CCA module can only model the reasoning mode of A$\rightarrow{*}\rightarrow$B, but cannot model the reasoning mode of A$\rightarrow{*}\leftarrow$B, A$\leftarrow{*}\leftarrow$B, and A$\leftarrow{*}\rightarrow$B. (3) Since most of the entity pairs are irrelevant in the document, the entity-pair matrix $M$ contains a lot of noise which may affect the reasoning ability of the model. Therefore, distinguishing related entity-pairs from unrelated entity-pairs and strengthening the interaction of the relationship feature vectors of related entity-pairs is the key to improving The reasoning ability of the model. To solve these problems, we have introduced the following methods:

The Transition module controls the dimensions of the new features generated by each layer of the Dense-CCNet model, which reduces the computational complexity of the model.

We use the Dense-CCNet to convert the original entity-pair matrix $M$ into a new context-enhanced entity-pair matrix $M^{\prime}$. Given an entity pair $(e_{s},e_{o})$, we first concatenate the two entity embedding ($h_{e_{s}}$, $h_{e_{o}}$) and new relation feature $M^{\prime}_{so}$, then we obtain the distribution of relationship via a bilinear function. Formally, we have:

For the loss function, we use adaptive-thresholding loss Zhou et al. (2021), which learns an adaptive threshold for each entity pair. The loss function is broken down into two parts as shown below:

where $TH$ is an introduced class to separate positive classes and negative classes: positive classes would have higher probabilities than $TH$, and negative classes would have lower probabilities than $TH$, $P_{D}$ and $N_{D}$ are the positive classes set and negative classes set in document D respectively.

Finally, our total loss function is defined as follows:

We evaluate our Dense-CCNet model on three public document-level RE datasets. The statistics of the datasets could be found in Appendix A.

• •

  DocRED Yao et al. (2019): The DocRED is a large-scale crowdsourced dataset for document-level RE, which was constructed from Wikipedia and Wikidata. The DocRED contains 3053 documents for training, 1000 for validating, and 1000 for the test. It involves 97 types of target relations in total, and each document approximately contains 26 entities on average.

• •

  CDR Li et al. (2016): The CDR is a biomedical dataset is constructed by using the PubMed abstracts, which aims to predict the binary interactions between Chemical and Disease concepts. The CDR contains only one relationship and consists of 1500 human-annotated documents in total. The CDR are equally split into training, development, and test sets.

• •

  GDA Wu et al. (2019): Similar to the CDR, the GDA is also a dataset in the biomedical domain, but is constructed by distant supervision from the MEDLINE abstracts. The GDA contains 29192 documents as the training set and 1000 as the test set. Since there is no development set, we follow Christopoulou et al. (2019) to divide the training set into two parts according to the ratio of 8:2 and use them as training set and development set respectively.

Our model was implemented based on PyTorch. We used cased BERT-base Devlin et al. (2018) as the encoder on DocRED and SciBERT-base Beltagy et al. (2019) on CDR and GDA. We set the number of layers of Dense-CCNet to 3. Our model is optimized with AdamW Loshchilov and Hutter (2017) with a linear warmup Goyal et al. (2017) for the first 6% steps followed by a linear decay to 0. All hyper-parameters are tuned on the development set, some of which are listed in Appendix B.

The experimental results of our model on DocRED are shown in Table 1. We followed Yao et al. (2019) and used $F_{1}$ and Ign$F_{1}$ as the evaluation metrics to evaluate the overall performance of the model. Ign$F_{1}$ denotes the $F1$ score excluding the relational facts that are shared by the training and dev/test sets. We compare the Dense-CCNet model with the following two types of models on the DocRED dataset :

• •

  Graph-based Models: these models first construct the document-graph from the document, and then perform inferences through GCN Kipf and Welling (2016) on the graph. We include GEDA Li et al. (2020), LSR Nan et al. (2020), GLRE Wang et al. (2020a), GAIN Zeng et al. (2020), and HeterGSAN Xu et al. (2021b) for comparison.

• •

  Transformer-based Models: These models directly use the pre-trained language models for document-level RE without graph structures. we compared $\rm BERT_{base}$ Wang et al. (2019), BERT-$\rm TS_{base}$ Wang et al. (2019), HIN-$\rm BERT_{base}$ Tang et al. (2020), $\rm CorefBERT_{base}$ Ye et al. (2020), Coref$\rm BERT_{base}$Ye et al. (2020), and ATLOP-$\rm BERT_{base}$ Zhou et al. (2021) with our model.

In addition, we also consider the DocuNet Zhang et al. (2021) model in the comparison, which formulates document-level RE as a semantic segmentation problem.

As shown in Table 1, our Dense-CCNet model achieved 62.74%$F_{1}$ and 62.55% $F_{1}$ in the training set and test set, which outperforms the state-of-the-art model with 0.91% $F_{1}$ and 0.69% $F_{1}$ respectively. Compared with the GAIN model that is the state-of-the-art model of graph-based methods, our model exceeds it by 1.52% $F_{1}$ on the dev set and 1.31% $F_{1}$ on the test set. This proves that the logical reasoning on the entity-pairs level is more effective than previous methods on mentions or entities level.

On the CDR and GDA data sets, we compared BRAN Verga et al. (2018), EoG Christopoulou et al. (2019), LSR Nan et al. (2020), DHG Zhang et al. (2020c), GLRE Wang et al. (2020a), ATLOP Zhou et al. (2021), and DocuNet Zhang et al. (2021) with our model. The experimental results on two biomedical datasets are shown in Table 2.

Our Dense-CCNet-$\rm SciBERT_{base}$ model obtained 77.06($\pm$0.71)% $F_{1}$ and 86.44($\pm$0.25)% $F_{1}$ on two data sets respectively, which is also the new state-of-the-art result. The Dense-CCNet-$\rm SciBERT_{base}$ improved the $F_{1}$ score by 0.76% and 1.14% on CDR and GDA compared with DocuNet-$\rm SciBERT_{base}$. These results demonstrate the strong applicability and generality of our approach in the biomedical field.

We conducted an ablation experiment to validate the effectiveness of different components of our Dense-CCNet model on the development set of the DocRED dataset. The results are listed in Table 3, where w/o Dense Connection replaces the Densely connected Criss-Cross Attention with the Recurrent Criss-Cross Attention (RCCA), w/o Expanding Attention uses standard Criss-Cross Attention and does not extend the field of attention, w/o Clustering Loss and w/o Attention Bias respectively removes the Clustering Loss and the Attention Bias from our model.

From Table 3, we can observe that the w/o Dense leads to a drop of 1.62% $F1$, which shows that the features of low-level inference are very helpful for relation extraction and the features of high-level inference may contain noises. The w/o Expanding Attention caused a performance drop of 0.83% $F1$, which indicates that document-level RE may include multiple inference modes and our model can effectively expand the reasoning mode of the Criss-Cross Attention through the Expanding the Field of Attention technology.

The w/o Clustering Loss module and w/o Attention Bias module led to performance degradation of 0.72% $F1$ and 1.14% $F1$ points respectively, which reflects that reducing noise (irrelevant entity-pairs) may be the key to further improving entity-pair level inference. We guess that the most ideal entity-pair-level reasoning method may be to only propagate information between related entity pairs.

In addition, we also introduced the ablation study of the number of layers of the Dense-CCNet, and the experimental results are shown in Table 4. When the number of layers increases from 2 to 3, our model can capture more multi-hop inference features, so the performance of the model is improved by 1.3% $F1$. However, when the number of layers is increased to 4, the performance drops slightly by 0.37% $F1$ points. The possible reasons is that the noise has a greater impact on the high-level feature or the model falls into over-fitting.

We followed GAIN Zeng et al. (2020) to select the same example and conduct a case study to further illustrate that our Dense-CCNet model can effectively capture the interdependence among entity-pairs and perform entity-pair-level logical reasoning compared with the baseline.

The experimental results are shown in Figure 3. Figure 3c demonstrates that our model has better logical reasoning ability than the baseline. Figure 3a shows that the entity pair (“Without Me", May 26, 2002) has more attention to the entity pairs (“Without Me", The Eminem Show) and (The Eminem Show, May 26,2002), which indicates that our model could capture the correlation these among entity-pairs.