模型解释新方向！浙大揭秘LLM隐层之间的知识流动！

作者：bhn

论文 ：https://arxiv.org/pdf/2405.17969 - NIPS2024
代码 ：https://github.com/zjunlp/KnowledgeCircuits

本文中：

• knowledge editing=知识编辑；
• machine unlearning=遗忘学习；
• detoxification=祛毒；
• knowledge circuits=知识回路；

背景

大量的参数赋予了大模型强大的能力的同时也带来了一些缺陷，如幻觉，不安全的行为。由于模型复杂的知识存储机制，这些问题的解决变得十分复杂。

目前的工作进行两个方向的尝试，一种是认为事实以键值对的形式存储在mlp中，在这个基础上人们使用 知识编辑，遗忘学习，祛毒 等方法对模型的mlp层修改，以缓解修复模型的缺陷。

另一种思路认为，知识不是单独的存储在某一区域的，而是由不同的组件共同构成的，本文中称为 知识回路。

这篇论文的工作分为四部分 发现Transformer架构中的知识回路： 解释 知识回路 的概念和知识回路是如何发现的。

知识回路揭示隐式神经知识表征 ：通过一个简单的案例分析 知识回路 是如何解释知识表征的。

知识回路阐明知识编辑的内部机制 ：通过 知识回路 解释对 ROME 和 FT 这两种 知识编辑 方法进行解释。

知识回路促进解释语言模型行为 ：通过 知识回路 解释模型的幻觉和上下文学习背后的原因。

本文的工作和结果分析

1. 发现Transformer架构中的知识回路

   知识回路 ：相对于知识编辑关注知识的存储区域，知识回路更关注信息的流动。将语言模型看作一个由组件（input，output，attention_head，mlp）为节点，连接组件的边（残差流），共同组成的一个计算图，信息在这些组件中流动。

   下图为模型回答“The official language of France is ”这个问题时，所经过回路，具体的回路发现在下面介绍。

图3: 在回路中，MLP14类似的点代表着第14层的MLP层；L18H14代表着第18层的第14个注意力头，点之间的褐色连线代表这他们之间的信息流动。
回路发现 ：简单的来讲，回路的发现通过消融节点之间的边（即参数置为0），判断这个边对于知识是否为关键边。对于上图的回路来说，一个事实三元组“(Franch, official language,  French)” ，让模型补全”The official language of France is “这句话，预测客体French，在这个过程中，消融组件之间的边，判断边对于这个事实的重要性，保留重要的边，从而构造出了关于这个事实的回路。（具体实现请查看代码）

2. 知识回路揭示隐式神经知识表征

图5:Target Entity at Last Position表示”French“这个词在”is“位置时输出logits的rank。Target Entity at Subject Position表示“French”这个词在“France”位置时输出logits的rank—越低rank越高。Prob. of Entity表示实体的可能性—越高可能性越大。
从图5可以看出，在MLP17层以后可以的看到Object Entity的可能性开始渐渐上升。让我们注意到图四a，连接MLP17的节点由L14H13、L14H7等。在图四b中可以发现，L14H13的输出logits更加关注“Language”这类关系型token，L14H7则是将信息从主体“France”流动到最后一个位置。

从这个简单的例子，该论文猜想，不同的注意力头起不同的作用，L14H13这类头更加关注关系信息，而L14H7则是在主体位置提取关系信息。而MLP层则是结合之前组件提供的信息，提高目标token的最高rank。

3. 知识回路阐明知识编辑的内部机制

对于同一个问题“Platform Controller Hub is created by Microsoft.”，原始模型在L15H3更加聚焦于一些firmware这类词，并且正确答案“Intel”的可能性一直为0。而对于ROME编辑过的模型从15层开始“Intel”的可能性渐渐上升。而对于FT编辑方法，“Intel”从第0层就开始上升。

从ROME图中Subject Position的曲线总是在Last Position上面，并且他们在一起下降，并且在15层开始慢慢上升（L15H3已经开始获取了真确的答案）。可以得出ROME方法会将正确的信息添加到主体token位置，并且由注意力头进行信息提取。