Self-distillation framework for document-level relation extraction in low-resource environments

Backbone model

The backbone model utilizes the “∗” embedding at the beginning and end of the mention to represent the entity associated with the mention embedding. To obtain the entity’s embedding for all mentions ${\left\{m_j^i\right\}}_{j=1}^{N_{e_i}}$ corresponding to e_i, logsumexp pooling (Jia, Wong & Poon, 2019) is employed as follows: ${\displaystyle {\mathrm{h}}_{e_i}=log{\sum }_{j=1}^{N_{e_i}}exp\left(h_{m_j^i}\right).}$

The pooling operation collects all the information referenced in the document, thereby generating an embedded representation of the entity. Experimental results demonstrate that the algorithm performs superiorly compared to average pooling.

Teacher Model

Various methods for enhancing entity representations can be incorporated into the teacher model to enhance overall model performance, simultaneously the performance of the backbone model can be improved by comparing the entity embedding with that of the teacher model. In comparison to the previous method RSMAN (Yu, Yang & Tian, 2022), we improve the performance of the backbone model in the teacher model by employing an attention mechanism to generate distinct entity representations for multiple candidate relations, as shown in Fig. 2.

For each candidate, the parameters of the random initialization can learn p_r as a representation of the relation. Then the representation of the dot product approach is used to calculate the relation p_r and every mention $m_j^i$ between semantic relevance as follows:

$s_{ij}^r=g\left(p_r,m_j^i\right),$

where g represents the function used to compute the semantic relevance between the two embeddings. This article employs the dot product to calculate the correlation; however, a multi-layer perceptron (MLP) can also serve this purpose.

Next, the softmax function is utilized to calculate the attention weights for all mentions corresponding to different relations, as shown below: ${\displaystyle {\alpha }_{ij}^r=\frac{exp\left(s_{ij}^r\right)}{\sum _{k=1}^{Q_i}exp\left(s_{ik}^r\right)}.}$ Once the attention for all entity mentions is obtained, the entity representation for a specific relation can be derived by taking the weighted sum of the attention weights across different relations. For a given relation r between entities $e_i^r$ as illustrated below: ${\displaystyle e_i^r=\sum _{j=1}^{Q_i}{\alpha }_{ij}^r{m}_j^i,}$ ${\displaystyle {\grave{h}}_{e_l}=\left[e_i^0,e_i^1,\dots ,e_i^r\right].}$

Backbone model

In the basic model, the entity embedding is mapped to the hidden state z using linear and nonlinear activation layers, and calculate the probability of relation r with the bilinear function and the sigmoid activation function, ${\displaystyle z_s=tanh\left(W_s{h}_{e_s}\right),}$ ${\displaystyle z_o=tanh\left(W_o{h}_{e_o}\right),}$ ${\displaystyle \mathrm{P}\left(r\mid e_s,e_o\right)=\sigma \left(z_s^{\top }{W}_r{z}_o+b_r\right),}$

where the model includes learnable weight matrices W_s ∈ ℝ^d×d, W_o ∈ ℝ^d×d, W_r ∈ ℝ^d×d. The representation of the same entity remains consistent across different entity pairs. To reduce the parameter count in the bilinear classifier, the embedded dimension is divided into k equally-sized groups, and apply the bilinear function and sigmoid activation function within each group. This calculation is used to determine the probability of relation r. ${\displaystyle z_s=\left[z_s^1;\dots ;z_s^k\right],}$ ${\displaystyle z_o=\left[z_o^1;\dots ;z_o^k\right],}$ ${\displaystyle g_i^{\left(s,o\right)}=\sigma \left({\sum }_{i=1}^k{z}_s^{i\top }{W}_r^i{z}_o^i+b_r\right),}$ ${\displaystyle g^{\left(s,o\right)}=\left[g_1^{\left(s,o\right)},g_2^{\left(s,o\right)},\dots ,g_d^{\left(s,o\right)}\right],}$

where the learnable parameters $W_r^i\in {\mathbb{R}}^{\frac{d}{k}\times \frac{d}{k}},\mathrm{i}=1\dots \mathrm{k}$ are utilized, g^(s,o) represents the probability of relation r for the entity pair (e_s, e_o). Embedding the groups in k different ways reduces the parameter count from d² to d²/k.

Teacher model

In the teacher model, we incorporate various methods to enhance the representation of entity pairs to improve the performance of the model. After the entity representations corresponding to specific relations are arranged in the order of relation types, we employ a linear layer and activation function to map these organized entities to the hidden state z, thereby obtaining the representations of the subject and object entities, respectively. The representations of the subject and object entities are as follows: ${\displaystyle z_S^{\prime }=tanh\left(W_s{h}_e^{\prime }\right),}$ ${\displaystyle z_o^{\prime }=tanh\left(W_o{h}_e^{\prime }\right),}$ ${\displaystyle g^{\prime \left(s,o\right)}=\sigma \left({\mathbf{z}}^{\prime }{s}^{\top }{W}_r{z^{\prime }}_o+b_r\right),}$

where the learnable parameters W_s ∈ ℝ^d×d, i = 1…k are utilized. g^(s,o) represents the probability of relation r for the entity pair (e_s, e_o). Similarly to the backbone model, we utilize k groups to reduce the parameter count. Next, we employ axial self-attention to enhance the neighborhood information along the axes for each entity pair (e_s, e_o). Although previous studies have used axial self-attention to enhance neighbourhood information for relationship classification, we argue that self-attention ignores relationships with samples other than axial. In self-attention, an attention mechanism computes the semantic relevance between the q vector and the k vector, and subsequently applies weights from this attention to the v vector to generate a new feature. In contrast, external attention (Guo et al., 2022) functions differently by calculating the relevance between the q vector and a pre-defined common k vector, and subsequently generating a feature map through multiplication of this attention with another externally learnable v vector. Taking this into account, we changed the self-attention in axial attention to an external attention, as shown in Fig. 3.

Axial attention is computed using external attention along the horizontal and vertical axes. A skip connection is added with each calculation along the axis. Given n × n entity list, for entity pair to (e_s, e_o) by aggregating axial entity pair element (e_s, e_i) and (e_i, e_o) information. The entity pair path traverses two hops between (e_s, e_i) and (e_i, e_o) relations. Furthermore, external attention is used to incorporate information beyond the two hops. Subsequently, this information is utilized to classify the relations between the entity pair (e_s, e_o), encompassing not only the information of the two entities but also incorporating multiple hops. For entity pair (e_s, e_o), ${\displaystyle r_w^{\left(s,o\right)}=r_h^{\left(s,o\right)}+{\sum }_{p\in 1\dots n}{softmax}_p\left(q_{\left(s,o\right)}^T{k}_{\left(s,p\right)}\right)v_{\left(s,p\right)},}$ ${\displaystyle r_h^{\left(s,o\right)}=g^{\left(s,o\right)}+{\sum }_{p\in 1\dots n}{softmax}_p\left(q_{\left(s,o\right)}^T{k}_{\left(p,o\right)}\right)v_{\left(p,o\right)},}$

let q_(i,j), k_(i,j), v_(i,j) be represented as $W_q^{\prime }{g}^{\prime \left(i,j\right)},W_k^{\prime }{g}^{\left(i,j\right)},W_v^{\prime }{g}^{\left(i,j\right)}$ respectively, where $W_q^{\prime }\in {\mathbb{R}}^{d\times d}$, $W_k^{\prime }\in {\mathbb{R}}^{d\times d}$, and $W_v^{\prime }\in {\mathbb{R}}^{d\times d}$ denote the learnable weight matrices of the model. In contrast to self-attention, $q_{\left(s,o\right)}^T{k}_{\left(s,p\right)}$ represents the semantic relevance between the element at position (s, o) and other elements in row p. The parameters k_(s,p) and v_(s,p), which are independent of the input, serve as the weight for the entire training dataset along the horizontal axis.

Experimental baseline

We employ RSMAN module to validate the efficacy of distillation in the entity enhancement module and compare the enhancement effects of RSMAN module in two different models. We also include two recent approaches for comparison: the DocuNet model, which captures entity pair dependencies using a U-shaped network inspired by semantic segmentation in the computer vision field, and the SIRE model (Zeng, Wu & Chang, 2021), which extracts intra-sentence and inter-sentence relations using sentence-level and document-level encoders, respectively.

Since our axial attention approach considers the relation between an entity and itself, unlike ATLOP backbone model, we need to ensure that the model input contains a square matrix of adjacency matrices comprising all possible entity pairs. To achieve this, we fill in the missing entity pairs in the adjacency matrices during the data construction phase. In order to validate the distillation effect after altering the input data distribution, we retrained ATLOP five times using BERT as the encoder and averaged the results after incorporating the missing entity pairs into the adjacency matrix. In line with the experimental setup, we trained ATLOP backbone model using the parameter settings provided in the article. The best results on the dev dataset were uploaded to CodaLab (https://codalab.lisn.upsaclay.fr/) to obtain the test dataset results, serving as a baseline for our experiments.

Evaluation metrics

F1 and Ign F1 are the evaluation measures used in this work. Ign F1 represents the F1 score for relation facts that are excluded from the training, validation, and test sets. It is calculated by taking away the public relations that these sets have in common. Furthermore, in order to obtain the official test results, we forward CodaLab the best results on the development dataset.

Main results

Recent studies have made extensive use of ATLOP as a generalized backbone as part of their methodology. To improve entity representation and entity pair representation by distillation, we have added RSMAN module and the axial attention module to ATLOP backbone model. The experimental results are denoted by a star (⋆). Based on the experimental findings, we can see that ATLOP’s performance is on par with SSAN (Xu et al., 2021), SIRE, and DocuNet. The efficacy of the entity and entity-pair representation enhancement modules based on the self-distillation method is first demonstrated by the observation that ATLOP with self-distillation improves the baseline F1 value by 0.56 (dev set) compared to the baseline F1 value during the model prediction phase.

Furthermore, we examine the lifting effect of SSAN and CorefBERT models simultaneously with the inclusion of RSMAN module. We can see that RSMAN module increase on the CorefBERT and SSAN models is approximately 0.7. Our distillation process, in contrast, successfully lowers the number of parameters without significantly sacrificing performance.

To confirm the effectiveness of the framework, we switched to SciBERT and conducted tests on the GDA dataset in addition to the comparison on the DocRED dataset. The experiment results are shown in Table 5 and show that our strategy is still effective on the GDA biological domain dataset.

In the ensuing analyses, we first evaluated the effects of entity representation and entity pair representation enhancement by using three common representational enhancements in ablation experiments. Then, we conducted qualitative results, validity, and parametric sensitivity analyses in the experimental analyses.

Entity representation enhanced module

We test the entity representation improvement module on ATLOP backbone model in order to verify its effectiveness. In particular, we examine the direct effects of RSMAN module on different baseline models and its influence on our model after distillation. RSMAN module shows improvements of 0.73 and 0.3 in CorefBERT and SSAN backbone models, respectively, as shown in Table 6. Additionally, we use the same DocRED dataset to incorporate RSMAN module into ATLOP and conduct distillation, which leads to a 0.39 improvement over ATLOP baseline F1 value. This accomplishment is largely in line with RSMAN module’s performance increase.

Moreover, we assess the performance of ATLOP when encoding the entity representation based only on the evidence sentences from the articles, using the original text as input for the teacher model during distillation, in order to offer more proof of the effectiveness of our entity representation enhancement module. Using only the evidence sentences at the prediction step, the findings shown in Table 7 show that our technique achieves a near approximation to the input produced by the original text. This result provides more evidence of the efficacy of our entity representation enhancement.

Entity pair representation enhanced module

We just concentrate on distilling the representation enhancement module in order to assess the efficacy of our entity pair representation enhancement module. We compare the effect of directly integrating axial attention into ATLOP backbone model with the effect obtained by distillation in our model with respect to the entity pair representation improvement module. When axial attention is given directly to ATLOP, as shown in Table 8, we find that the F1 value increases by 0.63. Furthermore, our model’s F1 impact following distillation shows a 0.38 improvement, which is close to the result of adding axial attention directly. This result demonstrates that our entity pair representation improvement module is successful.

Qualitative results

By reducing the number of parameters, our self-distillation model minimizes the amount of resources used during deployment. In addition to the parameters needed for ATLOP, the additional modules included to improve the entity and entity pair representations can be ignored via self-distillation in the model’s prediction phase. Moreover, different strategies might be used in place of these two modules.

Since the goal of this study is to utilize knowledge distillation to minimize the number of parameters in the model, we use RSMAN and axial attention modules as examples to show how well our knowledge distillation framework works. We determine the fraction of parameters that are occupied by the axial attention module and RSMAN module, respectively, and the statistics are shown in the Table 9.

Except BERT encoder, which only requires fine-tuning, we computed the parameter requirements for the backbone model, the entity representation enhancement module, and the entity pair representation enhancement module, as shown in Table 9. Specifically, the backbone model ATLOP necessitates approximately 40 million parameters. To enhance the entity representation and entity pair representation, we employ the RSMAN module and the axial attention module, respectively. RSMAN utilizes around 19 million parameters, while the self-attention-based axial attention module requires approximately 17.72 million parameters. On the other hand, the external attention-based axial attention module involves approximately 31.92 million parameters. Comparatively, the addition of parameters for entity and entity pair representation augmentation accounts for about 47% and 56% of the overall model, respectively.

We can see that the additional representation enhancement modules make up approximately half of the total model proportion. We can reduce the number of parameters by eliminating all of the additional representation enhancement modules after using the knowledge distillation framework to distill the model and save only the backbone model.

Efficiency analysis

Apart from calculating the required number of parameters, we conducted a comparative analysis of the model’s resource utilization during the deployment phase. As depicted in Table 10, saving only the essential parameters resulted in a 47.7% reduction in file size. Moreover, when freshly loading the model onto the graphics card, considering the GPU consumption by the model structure, the model size decreased by 30.9%. Experimental findings demonstrate that the incorporation of additional modules significantly contributes to memory and graphics memory consumption. However, our self-distillation framework helps mitigate these issues to some extent.

Parameter sensitivity analysis

In this experiment, we introduce two hyperparameters, α, and β, to ensure that the loss calculation methods remain within an appropriate range when computing the total loss. As a result, we will conduct parameter sensitivity analyses for both α and β.

By examining the variation of the loss function depicted in Fig. 4, it is evident that the values of α and β primarily impact the final loss function of the model during the initial stages. However, as the number of iterations increases, the losses for all modules of the model gradually converge towards a smaller value, rendering the influence of the hyperparameters α and β on the model’s losses negligible.

The results of our analysis are shown in Fig. 5. The left plot presents the impact of varying the α parameter while keeping the β parameter fixed at 0.3 on the model’s performance. Conversely, the other plot illustrates the effect of changing the β parameter while setting the α parameter to 0.3. The findings indicate that despite arbitrary adjustments to these parameters, the resulting loss values gradually converge to a smaller value, thereby indicating minimal impact on the final model performance.