论文解读:Knowledge Neurons in Pretrained Transformers

论文发表于自然语言处理顶会ACL-2022(原文链接)。本文引入知识神经元 (Knowledge Neuron) 的概念,初步研究了预训练Transformer中是如何存储事实知识的:

1、通过研究预训练的BERT,本文发现预训练语言模型的知识是独立存在于中间神经元激活中的

2、可以通过编辑FFN层权重来修改模型对相应知识的记忆,而不用进行微调。

阅读本文请同时参考原始论文图表。

思想

本文发现Transformer的知识在模型内部被其中的FFN模块的知识神经元所表示,所谓的知识神经元 (Knowledge Neurons) 即为输入通过第一个权重W_1再激活后的向量,如图1和图2所示。

因此提出一种特征归因方法来定位能表示某个知识的知识神经元,知识通过一个事实三元组来表示(h,r,t)。本文使用某个包含 h,r 的prompt让模型来预测 t ,搞成一个完形填空任务的形式。

知识归因

对于某个神经元w_i\^{(l)},当其为某个值\\hat{w}_i\^{(l)}时,根据某个待完形填空的prompt x,模型预测正确答案y\^\*的概率定义为文中式(4)。则这个神经元w_i\^{(l)}的归因分数定义为文中式(5),也就是一个集成梯度 (Integrated Gradients)。可以理解为这个神经元从无到有对模型输出正确答案的平均影响程度。这个积分文中实现为求和的形式,把求和区间定义为20.

知识神经元精炼

选择归因分数大于阈值t的神经元,得到神经元粗糙集作为某个prompt的答案的知识归因集合。一个prompt可能有多个归因分数大于阈值t的神经元被定位,但它们并不一定都是正确答案的知识神经元,也可能是其它信息。为了正确定位正确答案的知识神经元,文章使用多个有相同正确答案的prompt来得到多个知识神经元粗糙集,然后选择这些粗糙集中共享比例大于p的神经元作为正确答案的知识神经元,即3.3所谓的知识神经元精炼。

实验

实验设置

实验在BERT模型上进行,每个注意力模块隐藏层元素为768,FFN隐藏层元素为3072。

将归因分数阈值t设为最大归因分数的0.2倍,神经元共享比例阈值 p 初始化为0.7,然后在实验过程中逐步增加或减小这个值,直到知识神经元数量在[2,5]以内。

实验基于PARAREL数据集,PARAREL数据集由专家策划,包含来自T-REx数据集(ElSahar et al., 2018)的38个关系的各种提示模板。关系事实模板如表1所示。

实验结果

知识神经元识别结果

图3:发现大部分与事实相关的知识神经元分布在模型的顶层,也就是最后几层。

表2:关系事实识别出来的平均知识神经元数量,baseline是直接使用知识神经元本身\\overline{w}_i\^{(l)}作为归因结果。其中inrta-rel和inter-rel分别表示有相同关系的事实对和有不同关系的事实对,可以看出本文的归因方法能有效识别出相同关系事实对共享的知识神经元,并且不同关系的事实对没有贡献的知识神经元。

图4:通过抑制定位的知识神经元,即直接设置为0,发现本文归因方法有效使正确答案的预测概率降低,平均降低29.03%。而beseline几乎无效。

图5:通过增强定位的知识神经元,即将激活值加倍,发现本文归因方法有效使正确答案的预测概率提升,平均提升31.17%。而baseline几乎无效。

研究什么样的prompt能激活知识神经元

通过在Bing上爬取,构建了一个新的更广泛的数据集BINGREL,抓取数据如下:

1、限制特定头尾实体对的句子,每个对最多抓取10句。共抓取210217句。

2、仅限制头实体的句子,每个头实体最多抓取10句。共抓取266020句。

根据远程监督假设,第一类文本通常包含完整的关系事实,而第二类不一定包含。实验分为三组:

\\mathcal{T}_1:mask第一类文本的尾实体获得表达知识的prompt。

\\mathcal{T}_2:随机mask第二类文本部分单词,作为对照组。

\\mathcal{T}_3:使用随机采样的prompt,作为另一个对照组。

表4:展示了三个组的知识神经元定位结果,发现知识神经元更容易被表达知识的prompt即\\mathcal{T}_1激活。此外,由于知识神经元是通过PARAREL数据集提取的,而没见过BINGREL,所以展现了提取的知识神经元的泛化能力,它的确表达了某种知识。

表3:展示了在相应知识神经元上平均激活值(注意不是归因值)最大和最小的prompt,发现表达相应知识的prompt能使相应的知识神经元激活值更大。

事实更新与关系擦除

表6:将\对应的知识神经元选定的FFN的第二个线性权重的slots,修改为\对应的slots时的结果:将原始预测的t修改为t'的比例 (Change rate)、使t'的预测概率最大的成功率 (Success rate)、有相同关系r的其它prompt的PPL的提升(Delta Intra-rel. PPL)、不同关系的其它prompt的PPL的提升 (Delta Inter-rel. PPL)。后面两个表示了这样的修改操作对无关事实的影响程度,越小越好。其中修改操作就是5.1的Methods,但是为什么要用词嵌入来修改,还不太清楚。PPL (Perplexity)就是模型正确预测一个句子的概率的倒数的表征,越小越好。

表5:展示了将某个关系的知识神经元对应的FFN\^{(val)}擦除(直接设置为0向量)后,模型对有该关系的知识的预测的PPL大幅增加,而没有该关系的知识的预测的PPL增加不大。说明定位到的知识神经元的确能独立地表征相应的关系概念,并且删除后不会影响其它概念。定位方法如5.2 Methods所示。

未来工作

1、弄清楚知识神经元在多语言预训练模型中的工作机制。感觉可以用有相同事实的不同语言prompt来归因。(本文指出)

2、归因方法还可以进一步改进,毕竟文中的结果指标并不高。

3、大模型中的知识神经元机制研究。

4、知识神经元仅仅出现在FFN中吗?自注意力层是否也存在?两者是否有某种交互作用?