WISE：重新思考大语言模型的终身模型编辑与知识记忆机制

论文地址：https://arxiv.org/abs/2405.14768https://arxiv.org/abs/2405.14768

1. 概述

随着世界知识的不断变化，大语言模型（LLMs）需要及时更新，纠正其生成的虚假信息或错误响应。这种持续的知识更新被称为终身模型编辑。当前的模型编辑方法主要通过两种方式存储新知识：长期记忆（模型参数）和工作记忆（神经网络激活/表示的非参数化知识）。然而，这两种方式在终身编辑情境下存在局限性，无法同时实现可靠性、泛化性和局部性。为了解决这一问题，本文提出了一种新的模型编辑方法——WISE。该方法采用双参数内存机制，在预训练知识与更新知识之间实现了无缝衔接，通过一种知识分片机制来避免知识冲突。

大语言模型在训练时获取的知识是静态的，但世界知识是不断变化的。因此，部署后的LLMs可能生成包含幻觉、偏见或过时的响应。同时，由于重新训练或微调模型代价高昂，无法满足快速增长的知识需求。因此，终身模型编辑成为一种解决方案，通过不断地更新和注入知识来保持模型的准确性。然而，模型编辑面临一个关键问题：更新后的知识应存储在何处。

长期记忆编辑：直接修改模型参数存储新知识。这种方式虽然可以保持对新知识的记忆，但会导致与原有预训练知识发生冲突，影响局部性和可靠性。
工作记忆编辑：利用非参数化的检索机制存储新知识，不修改模型参数。这种方式虽然可以避免与原有知识的冲突，提升局部性，但难以泛化，无法应对多样化的查询。

在此背景下，如何在终身模型编辑中平衡可靠性、泛化性和局部性，成为模型编辑的一个核心难题。

为了解决上述问题，本文提出了WISE（双参数内存机制的终身模型编辑方法），通过结合长期记忆和工作记忆的优点，打破模型编辑中的“不可能三角”。

双参数内存机制：
- 主内存：存储预训练时的原始知识，保持长期记忆的泛化能力。
- 侧内存：专门用于存储和更新新知识，保证可靠性和局部性。
路由机制：为了确保编辑后的模型在不同情境下调用合适的知识，WISE设计了一个路由器。该路由器根据输入查询的不同，自动选择是使用主内存中的预训练知识，还是侧内存中的更新知识。这样就能避免对无关知识的干扰。
知识分片机制：为了应对不断增加的编辑请求，WISE引入了知识分片机制。每一组编辑知识会被存储在不同的参数子空间中，这些子空间是相互正交的，确保了不同编辑之间不会产生冲突。编辑完成后，这些分片会被合并为一个共享的侧内存，从而实现无冲突的终身编辑。

WISE的设计主要解决了现有模型编辑方法在终身学习中的不足。具体体现在以下几点：

可靠性：模型可以记住当前和之前的编辑内容，不会在多次编辑后遗忘新知识。
局部性：编辑只会影响到相关的知识，不会破坏与该编辑无关的预训练知识。
泛化性：模型不仅能记住具体的查询-目标对，还能理解编辑内容，能够应对不同形式的相同知识查询。

在实验中，WISE显著优于传统的基于长期记忆或工作记忆的模型编辑方法。其在多个LLM架构（如GPT、LLaMA、Mistral等）下的实验表明，WISE在问答任务、幻觉检测和分布外任务等终身模型编辑场景中，均在可靠性、泛化性和局部性三个指标上取得了更好的结果。

2. 方法

2.1 终身模型编辑

终身模型编辑问题，其核心目标是通过连续的编辑，使得大语言模型（LLMs）的输出能够符合人类的期望，同时保持模型之前的知识和能力。

2.2 重新思考终身模型编辑的内存设计

表1列出了当前主要模型编辑方法在内存类型和终身编辑能力上的对比。对比的关键维度包括：

长期记忆编辑（修改模型参数）
工作记忆编辑（检索时使用神经网络激活/表示的非参数化知识）
参数化知识与检索知识
是否支持终身编辑
可靠性、泛化性和局部性

表1中的方法对比总结如下：

FT-EWC、ROME、MEMIT、MEND：这些方法通过修改 LLM 模型参数来编辑长期记忆，但无法支持连续编辑，或者会对无关知识产生负面影响，导致局部性较差。
GRACE：基于工作记忆，通过检索知识避免了对无关知识的冲突，但由于检索机制只记忆查询，不理解其含义，泛化能力较差。

WISE 方法则通过结合长期记忆和工作记忆的优势，提供了一种更有效的终身模型编辑方案。它利用一个双参数内存机制，同时保持 LLM 的长期记忆（泛化能力）和工作记忆（可靠性和局部性），使其在终身模型编辑中表现出色。

2.3 WISE: 侧内存与知识分片、合并及路由

WISE 的侧内存包含两个关键组成部分：

侧内存设计：
- 侧内存：侧内存是 LLM 某一前馈神经网络（FFN）层的副本，用于存储编辑流。通过这种设计，避免直接修改主内存而可能带来的遗忘和副作用。
- 路由机制：为了决定使用主内存还是侧内存，设计了路由激活组件，来识别编辑的范围。在推理过程中，该组件决定是使用主内存还是侧内存来完成推理。
知识分片与合并：
- 知识在随机子空间中的分片：为了避免遗忘，将侧内存分为多个随机子空间，保证知识编辑的密度和分布性。
- 知识合并：利用模型合并技术，将不同的侧内存片段合并为一个共享的侧内存，避免知识丢失。

2.3.1 侧内存设计 (Side Memory Design)

（1）侧内存设计的基本原理：
在Transformer中，每层包含一个多头自注意力机制（MHA）和一个前馈神经网络（FFN），而FFN占据了模型中大量的参数。为避免直接修改主内存（模型预训练时学到的知识），WISE引入了侧内存，用来存储编辑过的知识。

（2）主内存与侧内存之间的路由 (Routing between Side Memories and Main Memory)

（3）基于边界的损失函数 (Margin-based Loss Function)

2.3.2 知识分片与合并 (Knowledge Sharding and Merging)

（1）知识密度问题 (Knowledge Density)

为了在终身模型编辑中有效存储更新的知识，作者引入了知识密度的概念，它类似于知识容量，用于描述在模型的参数中存储了多少知识。在这个背景下，存在以下两难问题：

知识密度过低：如果编辑次数较少或对整个内存进行微调，知识密度低，可能导致过拟合；
知识密度过高：如果编辑过于频繁，知识密度过高，导致已编辑的知识发生冲突，可能引发灾难性遗忘。

为解决此问题，作者提出了一种知识分片和合并机制，将侧内存编辑划分为多个子片段，存储在不同的参数子空间中，随后通过合并这些子空间，形成一个完整的侧内存。这样设计的好处在于避免了知识冲突，同时实现高效存储。

（2）随机子空间中的知识 (Knowledge in Random Memory Subspaces)

（3）知识合并 (Knowledge Merging)

在完成多次编辑后，多个子空间中的知识需要合并为一个共享的侧内存。由于不同的子空间通过随机掩码生成，这些子空间可能存在重叠部分和不相交部分。作者提出了如下定理来描述这些子空间重叠的情况：

（3）知识合并技术：Ties-Merge

Ties-Merge 的合并过程分为三步：

修剪：修剪每个任务向量中的冗余参数；
符号选择：为每个参数选择最合适的符号；
不相交合并：计算不相交子空间的参数均值，并将结果合并到一个统一的侧内存中。

通过 Ties-Merge，多个子空间中的知识能够有效合并，减少了子空间合并时的冲突。

（4）路由与检索多个侧内存 (Routing and Retrieving among Several Side Memories)

由于单个侧内存的知识容量有限，WISE 设计了一个多侧内存系统，能够产生多个侧内存并在推理过程中进行检索。检索过程通过激活评分路由机制（activation score routing）来实现，系统会根据不同的激活指示器分数，选择最合适的内存进行推理。该设计被称为WISE-Retrieve，允许模型应对更复杂的终身编辑场景。

3.实验

3.1 实验设置和评估指标

（1）数据集与模型

选择了几种流行的自回归大型语言模型（LLMs）进行实验，包括：
- LLaMA-2-7B
- Mistral-7B
- GPT-J-6B

论文使用了三个不同的数据集：

ZsRE（零样本关系抽取，用于问答任务）
SelfCheckGPT（用于修正语言模型生成的幻觉现象）
Temporal（用于评估编辑模型在分布外数据上的泛化能力）

表格3中给出了这些数据集的统计信息，以及编辑数据和评估时使用的无关数据。

（2）基线方法

实验中，比较了多种基线方法和WISE，包括：

FT-L：直接微调，使用了KL散度损失。
FT-EWC：基于弹性权重合并（EWC）的连续学习微调方法。
GPT风格编辑器：如ROME和MEMIT，用于批量编辑模型。
MEND：基于超网络的编辑器。
DEFER 和 GRACE：基于检索的记忆编辑方法。

（3）评估指标

每个编辑示例包含三个主要测试指标：

Reliability（可靠性）：模型编辑的成功率。
Generalization（泛化能力）：编辑后的模型在其他类似查询上的表现。
Locality（局部性）：编辑后的模型应保持无关数据的输出不变。

3.2 实验结果

（1）WISE的竞争性表现

WISE 在实验中展示了相对于基线模型的卓越性能，特别是在以下几个方面：

WISE 超越了现有方法，尤其是在长编辑序列任务中；
直接编辑长期记忆（如 ROME 和 MEMIT 等）会导致与预训练知识的冲突，导致局部性差；
使用检索工作记忆的方法（如 GRACE 和 DEFER 等）在泛化能力上表现不佳，难以适应多样化查询。

在 问答任务（QA setting） 中，编辑次数 T=1000 时，WISE 在 LLaMA 和 Mistral 模型上分别获得了 0.83 和 0.79 的平均分数，相较于最接近的竞争对手提高了 18% 和 11%。这说明了 WISE 在处理长序列编辑时具有良好的稳定性和有效的管理能力。

相比之下，尽管 MEND 和 ROME 在编辑初期表现良好，但随着编辑序列的扩展，它们的表现明显下降，尤其在局部性方面。直接编辑长期记忆的方式（如 MEMIT、FT-EWC 和 MEND）会显著破坏模型的知识结构，且在 T=100 或 1000 时表现出局部性的大幅下降。

（2）分布外泛化评估（Out-of-Distribution Evaluation）

理想的模型编辑方法应能够在复杂分布转换（distributional shift）中从公式化编辑例子泛化到自然文本。基于此，使用 Temporal 数据集 测试了分布外的泛化能力。WISE 在该数据集上取得了最佳表现，在 OOD Gen.（泛化能力） 和整体性能上表现出色，尤其是表5所展示的结果。