预训练Transformer中的知识神经元.

0. TL; DR

本文提出了“知识神经元”的概念,通过知识归因方法识别出表达特定事实知识的神经元。研究发现,这些知识神经元的激活与知识表达呈正相关。此外,通过抑制或增强知识神经元的激活,可以显著影响知识的表达。本文还尝试利用知识神经元在不进行微调的情况下编辑(更新和擦除)特定事实知识,为预训练Transformer的知识存储和编辑提供了新的视角。

1. 背景介绍

大规模预训练Transformer模型(如BERT、RoBERTa)在处理自然语言处理(NLP)任务时表现出色,尤其是在回忆训练语料库中的事实知识方面。这些模型通过在大规模语料库(如维基百科)上进行语言建模目标的训练,自然地充当了一个自由文本知识库。

尽管已有研究通过填空任务(cloze task)评估了预训练模型中的事实知识,但这些研究主要关注模型输出的整体准确性。本文旨在深入研究预训练Transformer中事实知识的存储方式,特别是通过引入“知识神经元”的概念来识别表达特定事实知识的神经元。

2. 方法介绍

作者将Transformer中的前馈网络(FFN)模块视为键值记忆(key-value memories),其中第一层线性层计算中间神经元的激活值,第二层线性层通过加权求和整合值向量。知识归因方法通过计算每个神经元对知识预测的贡献来识别知识神经元。

作者提出了一种基于集成梯度(Integrated Gradients)的知识归因方法来识别表达特定事实知识的神经元。这种方法的核心思想是通过计算每个神经元对知识预测的贡献来识别“知识神经元”。

给定一个输入提示 $x$ 和一个关系事实 $\langle h, r, t \rangle$,模型的输出 $P_x(\hat{w}^{(l)}_i)$ 定义为预训练模型预测正确答案 $y^*$ 的概率:

\[P_x(\hat{w}^{(l)}_i) = p(y^* | x, w^{(l)}_i = \hat{w}^{(l)}_i)\]

其中:

为了计算神经元 $w^{(l)}_i$ 的归因分数 $\text{Attr}(w^{(l)}_i)$,作者使用了集成梯度方法。具体来说,从 $w^{(l)}_i = 0$ 到 $w^{(l)}_i$ 的原始值,逐步计算梯度并进行积分:

\[\text{Attr}(w^{(l)}_i) = w^{(l)}_i \int_{0}^{1} \frac{\partial P_x(\alpha w^{(l)}_i)}{\partial w^{(l)}_i} \, d\alpha\]

其中:

直接计算连续积分是不可行的,因此作者使用了黎曼近似来近似计算积分:

\[\tilde{\text{Attr}}(w^{(l)}_i) = w^{(l)}_i \sum_{k=1}^{m} \frac{\partial P_x\left(\frac{k}{m} w^{(l)}_i\right)}{\partial w^{(l)}_i}\]

其中 $m$ 是近似步数,作者在实验中设置 $m = 20$。

通过上述归因分数,作者识别出归因分数大于某个阈值 $t$ 的神经元作为粗略的知识神经元集合。为了更准确地识别知识神经元,作者提出了一个细化策略。具体来说,对于同一个事实的不同提示,假设它们共享相同的“真正阳性”知识神经元,但不共享“假阳性”神经元。因此,通过保留多个提示中广泛共享的神经元,可以过滤掉“假阳性”神经元。具体步骤如下:

  1. 生成 $n$ 个多样化的提示。
  2. 对每个提示,计算每个神经元的归因分数。
  3. 对每个提示,保留归因分数大于阈值 $t$ 的神经元,得到粗略的知识神经元集合。
  4. 考虑所有粗略集合,保留被超过 $p\%$ 提示共享的知识神经元。

3. 实验分析

作者通过实验验证了知识神经元的激活与知识表达之间的正相关关系。具体来说,通过抑制(将激活值设为 0)或增强(将激活值加倍)知识神经元的激活,观察模型对正确答案的概率变化。实验结果表明,抑制知识神经元显著降低了正确答案的概率,而增强知识神经元则显著提高了正确答案的概率。

实验结果表明,大多数与事实相关的神经元分布在预训练Transformer的最顶层。这与先前的研究结果一致,表明知识神经元主要集中在模型的高层。

作者统计了知识神经元的数量和分布情况。实验结果表明,平均每个关系事实有4.13个知识神经元被识别出来。此外,作者还计算了不同关系事实对之间的知识神经元交集,发现同一关系的事实对共享更多的知识神经元,而不同关系的事实对几乎没有共享的知识神经元。

作者进一步研究了哪些提示能够激活知识神经元。实验结果表明,知识神经元更倾向于被表达对应事实知识的提示激活。作者通过从Bing搜索引擎爬取的文本数据构建了一个新的数据集BINGREL,用于更广泛的比较。实验结果表明,知识神经元的激活能够区分表达知识的提示和不表达知识的提示。

作者尝试利用知识神经元在不进行微调的情况下编辑预训练Transformer中的特定事实知识。具体来说,作者进行了两个案例研究:更新事实和擦除关系。