【自然语言处理】【可解释性】NKB：用于预训练Transformers的神经知识银行

NKB：用于预训练Transformers的神经知识银行 《Neural Knowledge Bank for Pretrained Transformers》

论文地址：2208.00399.pdf (arxiv.org)

相关博客
【自然语言处理】【可解释性】NKB：用于预训练Transformers的神经知识银行
【自然语言处理】【可解释性】自注意力归因：解释Transformer内部的信息交互
【深度学习】【积分梯度】深度网络的公理归因(Axiomatic Attribution for Deep Networks)
【自然语言处理】【可解释性】Perturbed Masking:分析和解释BERT的无参数探针
【机器学习】【可解释性】LIME
【自然语言处理】【聚类】TELL：可解释神经聚类
【自然语言处理】【Prompt】语言模型即知识库(Language Models as Knowledge Bases)
【自然语言处理】【Prompt】P-tuning

一、简介

​ 近些年，大规模的预训练$\text{Transformers}$极大的改善了自然语言处理。它们通常会在一个大规模语料上训练，其中包含海量的事实知识。当面对一些像closed-book问题这样的知识敏感下游任务，记住事实知识的能力将至关重要。Petroni et al(2019)展示了预训练$\text{Transformers}$能够以zero-shot的方式召回出现在训练数据中的事实知识。Roberts,Raffel and Shazeer(2020)也证明了在微调后，$\text{T5}$能够在不访问外部知识上下文的情况下回答开放域的问题。即便如此，预训练模型存储事实知识的能力仍然有限。考虑到这一现状，作者设计了一种可解释方法方法来为预训练$\text{Transformers}$引入外部的事实知识。

​ Geva et al.(2020)指出$\text{Transformers}$中的$\text{Feed-Forward Networks}$类似于$\text{key-value memories}$。如上图所示，在$\text{FFN}$总，把第一个线性层看成是一系列key并且第二个线性层看作是对应的value。$\text{FFN}$的输入hidden state被送入第一个线性层，并激活一组神经元。然后，将这些激活神经元作为权重，第二个线性层通过加权求和来集成对应的值向量。在这个角度的基础上，Dai et al.(2022)进一步发现预训练$\text{Transformers}$中的$\text{FFNs}$以类似内存的方式存储事实知识。

​ 受这些发现的启发，作者设计了神经知识银行$\text{(Neural Knowledge Bank,NKB)}$和一个知识映射策略来向预训练$\text{Transformers}$注入额外的事实知识。$\text{NKB}$是一个类似$\text{FFN}$的模块，用于原始的$\text{FFN}$后拼接额外的知识存储槽。为了注入事实知识，先从$\text{Wikipedia}$中获得知识语料，然后使用$\text{Salient Span Masking(SSM)}$预训练目标函数来预训练$\text{NKB}$。在知识注入过程中，固定原始的预训练模型来避免影响其通用语言建模的能力。对于下游任务，可以直接微调整个模型。$\text{NKB}$的优点反应在三个方面：(1) $\text{NKB}$的知识注入过程是独立于原始的预训练模型，所以引入额外的知识并不会降低模型的通用语言建模能力。(2) $\text{NKB}$类似内存的架构具有高的解释性，并且能够以人类可读的方法介绍其keys和values的含义。(3) $\text{NKB}$的$\text{key-value}$架构高度灵活，并且能够通过修改value值来更新$\text{NKB}$的知识。

​ 在三个不同领域closed-book问题数据集中，作者发现挂载$\text{NKB}$能够提升$\text{T5}$的表现，特别是在生物医学领域，$\text{T5}$的预训练语料覆盖的比较少。通过摘要和机器翻译两个表示任务，证明了挂载$\text{NKB}$并不会使原始$\text{T5}$模型降低通用语言建模能力。此外，通过分析$\text{NKB}$揭露了其工作机制，并以人类可读的方式来展现key和value的含义。最后，通过修改值向量来更新存储的事实知识来展示$\text{NKB}$的灵活性。

​ 本文的贡献如下：

• 提出了$\text{NKB}$的想法，包括架构和知识注入策略，来向预训练$\text{Transformer}$注入额外的事实知识。
• 验证$\text{NKB}$能够提高$\text{T5}$在closed-book问答上的效果，并且保存$\text{T5}$的通用语言建模能力。
• 分析了$\text{NKB}$的可解释性，并且以人类可读的方式揭示了key和value的含义。
• 通过直接修改value向量来更新存储的知识来展示$\text{NKB}$的灵活性。

二、$\text{Transformer}$

​ 近期，$\text{Transformer}$已经成为$\text{NLP}$中最流行和有效的架构。以标准的$\text{Transformer}$编码器为例，其堆叠了$L$个相同的$\text{Transformer}$层。在每个$\text{Transformer}$层都有两个主要的模块：自注意力模块$\text{(self-attention,SelfAtt)}$和前馈网络$\text{(feed-forward network,FFN)}$。对于一个长度为$len$的输入序列，令$X\in {\mathbb{R}}^{len\times d}$表示一个$\text{Transformer}$层的input hidden states，两个模块的形式化为：
$$Q_h=X{W}_h^Q,K_h=X{W}_h^K,V_h=X{W}_h^V\text{\hspace{1em}(1)}$$

$${\text{SelfAtt}}_h(X)=\text{Softmax}(Q_h{K}_h^T)V_h\text{\hspace{1em}(2)}$$

$$\text{FFN}(H)=\text{ActFunc}(H{W}_1^T)W_2\text{\hspace{1em}(3)}$$

其中，$W_h^Q,W_h^K,W_h^V\in {\mathbb{R}}^{d\times \frac{d}{n}},W_1,W_2\in {\mathbb{R}}^{4d\times d}$是参数矩阵，${\text{SelfAtt}}_h(X)$是计算$n$个注意力头中第$h$的注意力，$H\in {\mathbb{R}}^{len\times d}$表示自注意力模块输出的hidden states，其通过将所有注意力头的平均投影来计算，并且$\text{ActFunc}$像$\text{GELU}$这样的激活函数。为了简单，这里忽略了自注意力的缩放因子和偏差项。

​ 比较等式(2)和等式(3)，可以发现$\text{FFN}(\cdot )$的计算形式与${\text{SelfAtt}}_h(\cdot )$的几乎相同，除了具有不同的非线性激活函数。因此，将$\text{FFN}$看作是具有query-key-value机制的模块是合理的。具体来说，$\text{FFN}$的输入$H$作为queries，第一个线性层的权重$W_1$是keys，第二个线性层的权重$W_2$是values。

三、方法

​ 遵循Geva et al.2020,2022;Dai et al. 2022，本文也将$\text{Transformer}$中的$\text{FFN}$看作是key-value存储。基于这个视角，作者设计了一个神经知识银行$\text{(Neural Knowledge Bank,NKB)}$和一个知识注入策略来向预训练$\text{Transformers}$中注入额外的事实知识。在本小节中，将会引入将$\text{FFN}$作为key-value存储的视角，$\text{NKB}$的架构和知识注入方法。

1. $\text{FFN}$的$\text{Key-value}$存储视角

​ 类似于Geva et al.(2020);Dai et al(2022)，将$\text{FFN}$形式化为$\text{key-value}$存储。正如上图所示，将一个token的输入hidden state $\text{h}\in {\mathbb{R}}^d$看作是query，两个线性层的参数矩阵$W_1,W_2\in {\mathbb{R}}^{4d\times d}$看作是$4d$的key和$4d$的value，其中每个key或者value是一个$d$维向量。首先，$\text{h}$别被输入至第一个线性层。对于$W_1$中的每个key向量，通过内积计算一个标量分数$s_i$：
$$s_i={\text{h}}^{\top }{W}_1[i,:]\text{\hspace{1em}(4)}$$
其中，$[\cdot ,\cdot ]$表示矩阵的切片操作。在激活函数之后，这些分数构成了中间隐藏层$\text{FFN}$的${\text{h}}^{(inter)}\in {\mathbb{R}}^{4d}$：
$$w_i=\text{ActFunc}(s_i)\text{\hspace{1em}(5)}$$

$${\text{h}}^{(inter)}=[w_1;w_2;\dots ,w_{4d}]\text{\hspace{1em}(6)}$$

其中，$[\cdot ;\cdot ]$表示拼接。最终，将这些激活的神经元作为权重，并通过获得$\text{FFN}$的${\text{h}}^{(output)}\in {\mathbb{R}}^d$来集成$W_2$中的value向量。
$${\text{h}}^{(output)}=\sum _{i=1}^{4d}{w}_i{W}_2[i,:]\text{\hspace{1em}(7)}$$

2. 神经知识银行($\text{Neural Knowledge Bank}$)

​ Dai et al.(2022)发现$\text{Transformers}$中的$\text{FFNs}$能够存储事实知识，隐藏作者设计的$\text{NKB}$架构与$\text{FFN}$相同。具体来说，对于一个$\text{Transformers}$层，分配两个新矩阵$W_1^{\prime },W_2^{\prime }\in {\mathbb{R}}^{d^{\prime }\times d}$作为额外的key和value，其中$d^{\prime }$是控制$\text{NKB}$容量的超参数。如上图所示，通过将$W_1^{\prime }$和$W_2^{\prime }$拼接在原始$\text{FFN}$的$W_1$和$W_2$之后来将$\text{NKB}$挂载至$\text{Transformer}$层。

​ 对于$\text{NKB}$，$\text{FFN}$的中间hidden state ${\text{h}}^{(inter)}$将会被扩展至$(4d+d^{\prime })$维：
$$s_i^{\prime }={\text{h}}^{\top }{W}_1^{\prime }[i,;]w_i^{\prime }=\text{ActFunc}(s_i^{\prime })\text{\hspace{1em}(8)}$$

$${\text{h}}^{(inter)}=[w_1;w_2;\dots ,;w_{4d};w_1^{\prime };w_2^{\prime };\dots ,w_{d^{\prime }}^{\prime }]\text{\hspace{1em}(9)}$$

其中，$w_i^{\prime }$是第$i$个额外存储槽的权重。最终，将新的${\text{h}}^{(inter)}$作为权重，value向量$W_2$和$W_2^{\prime }$被如下方式集成：
$${\text{h}}^{(output)}=\sum _{i=1}^{4d}{w}_i{W}_2[i,:]+\sum _{i=1}^{d^{\prime }}{w}_i^{\prime }{W}_2^{\prime }[i,:]\text{\hspace{1em}(10)}$$
作为$\text{FFN}$的简单扩展，$\text{NKB}$非常容易使用和实现。更重要的是，$\text{NKB}$这种类似于存储的架构高度可解释，并且能以人类可理解的方式解释key和value的含义。此外，key-value架构高度灵活，并且能够通过更新value向量来轻易执行知识更新。

3. 知识注入

​ 为了能够引入额外的事实知识，使用$\text{Salient Span Masking(SSM)}$目标函数来预训练$\text{NKB}$。具体来说，先使用$\text{DrQA}$的文档检索器来从$\text{Wikipedia}$中获得包含额外知识的语料。然后，使用$\text{spaCy}$中的实体识别器来识别每个文本片段中的salient spans(例如命名实体和日期)。最后，随着遮蔽salient spans并训练$\text{NKB}$的参数来重构被遮蔽的salient span。

​ 在知识注入的过程中，固定原始预训练模型的参数并仅更新$\text{NKB}$的参数。相较于先前的工作为知识注入更新预训练参数，本文的训练策略可以避免预训练模型的通用语言建模能力退化。此外，使用这个训练策略，新的事实知识被明确注入至$\text{NKB}$，因此也更容易定位和分析新引入的知识。在知识注入后，可以直接微调整个模型来用于下游任务。

四、实验

• 添加知识为问答任务带来的效果

• 添加知识是否会降低模型的通用语言建模能力

五、$\text{NKB}$的解释性

1. value向量存储实体

​ Dai et al.(2022)发现$\text{FFNs}$能够存储事实知识，Geva et al.(2022)表示$\text{FFNs}$中的value向量通常会对应人类可读的概念。受这些启发，作者在输出词典空间上分析了$\text{NKB}$的值向量，发现大多数的值向量都存储了具体的实体。

• 方法

  从在$\text{WebQuestions}$上微调的最好模型${\text{T5}}_{base}+\text{NKB}$随机采样$\text{NKB}$的值向量。对于第$i$个值向量${\text{v}}_i\in {\mathbb{R}}^d$，首先将其投影至输出词表空间并得到一个概率分布${\text{p}}_i\in {\mathbb{R}}^{N_{vocab}}$：
  $${\text{p}}_i=\text{Softmax}(E{\text{v}}_i)\text{\hspace{1em}(11)}$$
  其中，$E\in {\mathbb{R}}^{N_{vocab}\times d}$是${\text{T5}}_{base}$中的输出嵌入矩阵。然后，对于每个采样的值向量${\text{v}}_i$，手动检测${\text{p}}_i$中最高概率的token，并将其分类至6个实体类别：Person, Place, Organization, Date, Other或者Non-entity。为了保证人工标注的置信度，忽略不在英语中或者不完成单词的token。最后，作者标注了100个验证value向量。

• 结果

对于每个实体类别，上表会展示一个值向量示例和评分最高的token。此外，还在上图中展示了所有类别的评分最高token的分布。发现85%采用的值向量都是对应具体的实体。值向量中的实体信息将会被集成至$\text{FFN}$的输出hidden state ${\text{h}}^{(output)}$并有助于最终答案的生成。

2. key向量捕获输入模式

​ 相比于$\text{NKB}$的value向量存储输入独立的实体，$\text{NKB}$的key向量则是输入敏感的。它们负责根据输入模式来决定何时激活存储槽。在本小节，通过分析何时存储槽被激活，可以揭示触发key向量的模式。

• 方法

  具体来说，将$\text{WebQuestions}$验证集中的每个问题输入至模型中。当模型正在为一个问题生成第一个答案token时，记录$\text{NKB}$的权重。在记录完所有问题的$\text{NKB}$权重后，随机采样100个key向量进行分析。对于第$i$个key向量${\text{k}}_i\in {\mathbb{R}}^d$，选择出5个触发该key向量最多的问题(即拥有最高的$\text{NKB}$权重)。最后，人工检测这5个问题并尝试确定出人类可读的语义模型。

• 结果

从人工的标注中发现78%采样的key都具有与top触发问题共同的语义模式。在上表中展示了key向量和对应的top触发问题的例子。将关于$\text{NKB}$的key和value发现放在一起，可以总结出$\text{NKB}$的工作机制：key向量根据输入模式决定是否激活存储槽，激活的存储槽value向量被集成来用于生成最终的答案。

六、$\text{NKB}$的知识更新

​ 知识更新能力对于避免模型产生错误和过时的结果是有意义的。利用$\text{NKB}$灵活的key-value架构，作者提出了一个简单的方法，通过在$\text{NKB}$上执行一个小的知识手术来直接将一个问题的答案更新至另一个问题。

• 数据

  从$\text{WebQuestions}$的验证集合中选择样本用于知识更新。为了保证目标答案的可验证性，选择那些模型预测错误的问题，目标是将错误回答更正。在本文中，我们仅考虑单个token作为答案的问题并预测答案。最后，为后续的知识更新实验保留了67个样本。

• 方法

  受Dai et al.(2022)的启发，在$\text{NKB}$上执行一个小的知识更新手术来删除原始答案的信息，并引入目标答案的新信息。具体来说，对于一个问题，先选择具有最高$\text{NKB}$权重的存储槽。然后，在$\text{NKB}$执行下面的操作：
  $${\text{v}}_t+=\lambda ({\text{e}}_{tgt}-{\text{e}}_{ori})\text{\hspace{1em}(12)}$$
  其中，$t$是最高权重存储槽的索引，$\lambda$是超参数，${\text{e}}_{ori}$和${\text{e}}_{tgt}$表示原始答案和目标答案的输出词嵌入。直觉上，只需要修改value向量，知识手术倾向于抑制原始答案的预测概率，并鼓励模型生成目标答案。

• 评估

  使用两个指标来评估知识更新手术的效果。success rate表示答案成功更新至预期答案的比例。使用destruction rate来衡量知识手术对其他样本的影响。对于每个知识手术，随机采样5个问题，并且destruction rate表示知识手术改变问题答案的比例。

• 结果
  

  上表中展示了知识更新的结果。随着$\lambda$的增大，success rate和destruction rate同时增大。特别地，当$\lambda =0.07$，可以实现一个很高的success rate(98.5%)，并保存一个可容忍的destruction rate(2.7%)。