bert关系抽取论文源码之REDN：Downstream Model Design of Pre-trained Language Model for Relation Extraction Task

前言

REDN是一个使用bert预训练模型进行关系抽取的衍生模型，该模型设计了一种具有新型损失函数的网络结构，并且仅用bert-base-uncased模型便在三个主流关系抽取数据集SemEval 2010 Task 8 , NYT , WebNLG 上达到了sota的效果。

该模型的最大创新之处在于用矩阵而非一维向量表示实体之间的关系，并设计了相应的损失函数。文章认为关系是由首尾实体向量之间的相关性决定的，那么关系自然应该表示为一个矩阵而不是一维向量。这样在预测它们之间的关系时，就可以使用更多信息，如实体在文中的位置等。

本文将结合论文以及pytorch源码来解读该模型，本模型基于清华的开源关系抽取库 OpenNRE 实现，源码与论文详见参考资料。

模型架构

模型整体架构如图所示：

REDN模型由三个主要部分构成，一个部分是bert的encoder，用于提取文本的向量表示，同时从获得的文本向量中提取出关系的矩阵表示。第二个部分是Relation Computing Layer，用于计算头实体和尾实体之间的相关性。第三个部分是Loss Calculation。接下来我们将从这三个方面详细解读这个模型。

1.Encoder

在本部分，模型主要的目的是计算头实体和尾实体的矩阵表示Eb和Ep。这里回顾一下 bert encoder 的结构，主要是由 12 层transformer encoder 堆叠形成。通常情况下，我们会使用 bert encoder 中最后一层 transformer encoder 输出的大小为（batch_size, sequence_len, hidden_size）的隐藏向量，来作为文本的向量表示。而文章中提到，研究表明 bert encoder 中每一层输出的隐藏向量都可以作为文本向量表示。而在本文中，则使用 bert encoder 中倒数第二层的输出向量作为初始向量Ew。为什么不使用最后一层的输出向量呢？接下来我们将会进行说明。

为了获得尾实体的矩阵表示，文章使用bert encoder中最后一层的transformer encoder来处理Ew，公式如下：

说白了其实就是使用了bert encoder中最后一层的输出向量，这也解答了刚才产生的疑问。

文章中解释这样的处理可以让Ep不仅融合Ew的信息，同时也融合来自其他token的位置信息。

Such an operation is applied so that the embedding of every token in Ep will, in addition to Ew , fuse some information from tokens in other positions.

为了更好的获得模型的信息，文章将bert的常用于文本分类的特殊符号[CLS]的向量表示Ea与Ew相加，得到头实体的向量表示Eb。

这里Ea的shape是（batch_size, hidden_size），Eb的shape是（batch_size, sequene_len, hidden_size），因此实行的是广播运算，使头实体中每个词向量都更好的获取模型信息。

这里附上这部分的源码：

class BERTHiddenStateEncoder(BERTEncoder):
    def __init__(self, pretrain_path, blank_padding=True):
        super().__init__(80, pretrain_path, blank_padding)
        self.bert = BertModel.from_pretrained(pretrain_path, output_hidden_states=True,output_attentions=True)

    def forward(self, token, att_mask):
        """
        Args:
            token: (B, L), index of tokens
            att_mask: (B, L), attention mask (1 for contents and 0 for padding)
        Return:
            (B, H), representations for sentences
        """
        _, x, hs ,atts= self.bert(token, attention_mask=att_mask)
        return x, hs, atts

#这里的self.sentence_encoder就是BERTHiddenStateEncoder
#self.subject_1为False
#hs是一个大小为13的tuple，里面是一个bert经过embedding后得到的输出向量
#以及bert encoder所有12个transformer encoder层的输出向量

sentence_encoder = encoder.BERTHiddenStateEncoder(pretrain_path=bert_path)
self.sentence_encoder =sentence_encoder 
rep, hs, atts = self.sentence_encoder(token, att_mask)  # (B, H)
if self.subject_1:
    subject_output = hs[-1]  # BS * SL * HS
else:
    subject_output = hs[-2]  # BS * SL * HS
if self.use_cls:
    subject_output = subject_output + rep.view(-1, 1, rep.shape[-1])

2.Relation Computing Layer

这一部分的主要目的是为了得到相关性矩阵P。文章使用了非对称核内积函数来计算头、尾实体的相关性，公式如下：

这里Whi，Wti分别是头实体，尾实体的关系嵌入矩阵。它们是在训练过程中学习到的参数。而这里的Si矩阵每个关系都有一个，shape为(sequnce_len, sequnce_len)，代表着文本中各个token之间具有第i个关系的可能性。

为了加深理解，下面我们来看一下这部分的源码：

#这里的num_heads就是关系的数量
#q是hs[-1], k是subject_output

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, input_size, output_size, num_heads, output_attentions=False):
        super(MultiHeadAttention, self).__init__()
        self.output_attentions = output_attentions
        self.num_heads = num_heads
        self.d_model_size = input_size

        self.depth = int(output_size / self.num_heads)

        self.Wq = torch.nn.Linear(input_size, output_size)
        self.Wk = torch.nn.Linear(input_size, output_size)

    def split_into_heads(self, x, batch_size):
        x = x.reshape(batch_size, -1, self.num_heads, self.depth)  # BS * SL * NH * H
        return x.permute([0, 2, 1, 3])  # BS * NH * SL * H

    def forward(self, k, q):  # BS * SL * HS
        batch_size = q.shape[0]

        q = self.Wq(q)  # BS * SL * OUT
        k = self.Wk(k)  # BS * SL * OUT

        # q = F.dropout(q, 0.8, training=self.training)
        # k = F.dropout(k, 0.8, training=self.training)

        q = self.split_into_heads(q, batch_size)  # BS * NH * SL * H
        k = self.split_into_heads(k, batch_size)  # BS * NH * SL * H

        attn_score = torch.matmul(q, k.permute(0, 1, 3, 2))
        attn_score = attn_score / np.sqrt(k.shape[-1])

        # scaled_attention = output[0].permute([0, 2, 1, 3])
        # attn = output[1]
        # original_size_attention = scaled_attention.reshape(batch_size, -1, self.d_model_size)
        # output = self.dense(original_size_attention)

        return attn_score

#这里的hs就是encoder代码中的bert encoder输出向量的一个集合
#hs[-1]就是Ep，即bert encoder最后的输出向量
#subjuect_out就是Eb矩阵。

self.attn_score = MultiHeadAttention(input_size=hidden_size,
                                     output_size=num_class * hidden_size,
                                     num_heads=num_class)

score = self.attn_score(hs[-1], subject_output)  # BS * NTL * SL * SL
score = score.sigmoid()

可以看到，Ep, Eb经过一个全链接也就是文中与W矩阵做矩阵乘积之后，变成了shape为（batch_size, sequence_len, num_heads*hidden_size）的矩阵，再经过split_into_heads划分后，得到了了shape为
（batch_size, num_heads, sequence_len, hidden_size）的关系嵌入矩阵，即将原来的文本嵌入矩阵（batch_size, sequence_len, hidden_size）嵌入到每个关系中，使得对于每个关系，都有一个对应的和原来文本嵌入shape相同的嵌入矩阵。

接下来, 将q，k矩阵即Ep, Eb矩阵相乘得到所有的S矩阵。最后在对于得分attn_score也就是S矩阵除以一个系数加以规范，目的是使梯度更加稳定。这个在Attention is all you need论文的transformer结构中也有类似的处理。

将S矩阵激活之后，就可以得到相关矩阵P，表示每个token之间具有第i个关系的概率值，也就是本部分源码的最后一行。公式如下：

3.Loss Calculation

由于P代表的是每个tokenn之间的相关性，而非实体，因此文章构造了一个mask矩阵M，用于将除了实体以外的token遮蔽。公式如下：

S表示实体对的集合，m,n是M矩阵的下标，Mmn表示M矩阵位置是（m，n）元素，Bx, Ex表示一个实体的开始和结尾。该公式表示，当头实体或尾实体的位置包含（m，n）时，M该位置上的元素为1。

对于label的处理也类似，公式如下：

Yi表示具有第i个关系的实体对集合。m，n，BX，Ex与上文相同。这样的处理也解决了当一对实体具有多种关系的情况。

对与第i个关系的损失，计算公式如下：

这里的 * 是 Hadamard product ，即按位相乘。刚才提到P矩阵表示每个token之间的相关性，经过掩码矩阵M的处理之后，就可以得到实体之间具有第i个关系的概率。

最后的损失是所有关系损失的加和，公式如下：

损失函数的部分的源码如下：

class PARALoss(nn.Module):
    """
    Softmax classifier for sentence-level relation extraction.
    """

    def __init__(self):
        """
        Args:
            sentence_encoder: encoder for sentences
            num_class: number of classes
            id2rel: dictionary of id -> relation name mapping
        """
        super().__init__()

    def forward(self, score, predicate_one_hot_labels):
        entity_mask = predicate_one_hot_labels.sum(dim=1, keepdim=True).repeat_interleave(score.shape[1], dim=1)
        entity_mask = (entity_mask > 0).float()

        entity_sum = (entity_mask != 0).sum(dim=(2, 3)).float()  # BS, NL

        loss = ((F.binary_cross_entropy(score, predicate_one_hot_labels, reduction="none") * entity_mask).sum(dim=(2, 3))
                / entity_sum).mean()
        if loss.item() < 0:
            print("debug")
        return loss

结语

REDN模型是一个优秀的关系抽取模型，其使用矩阵表示关系嵌入的思想十分巧妙，同时也在三个数据集上达到了sota效果。

参考资料

论文：https://arxiv.org/abs/2004.03786

源码：https://github.com/slczgwh/REDN