TokenMem: 面向冻结大语言模型的忠实知识注入

# TokenMem：面向冻结大语言模型的知识忠实注入方法

来源：https://arxiv.org/html/2607.22625
程章宇¹ 陈阳正¹ 卢泽宁¹ 何英如¹ 黄宇彤¹ 张一鸣² 许悦¹ 靳占鹏¹

¹华南理工大学 ²中国科学技术大学

###### 摘要

检索增强生成（RAG）通过外部知识增强了大型语言模型（LLMs）的能力，但存在*知识冲突*问题：当检索到的信息与参数记忆相矛盾时，共享的自注意力路径会产生不可预测的输出。我们提出了**TokenMem**，一种轻量级记忆系统，通过专用的交叉注意力通道向冻结的LLM注入知识，从而避免了与残差流中参数记忆的竞争。**TokenMem**仅训练一个轻量级门控适配器（约3-7M参数），采用两阶段课程：首先学习通用知识利用，然后在反事实知识条件下强化忠实遵从。在涵盖三个系列（Qwen3-4B/8B/14B、LLaMA-3.1-8B、OLMo-3-7B）共五个模型的受控实验中，**TokenMem**在反事实基准测试上达到了69-70%的知识遵从率（KC），而标准RAG仅为20-52%——差距最大达49个百分点。消融研究表明，两阶段课程至关重要：移除第二阶段将导致KC骤降至接近零。机制分析显示，门控适配器在没有显式监督的情况下学习到了一种冲突感知、逐层特定的注入策略。代码可在我们的匿名仓库获取(https://anonymous.4open.science/r/TokenMem0-D35B/README.md)。

## 1 引言

大型语言模型（LLMs）的参数知识不可避免地存在不完整性，并随着世界发展而变得过时[16，33，18]。在金融合规、临床决策支持或法律咨询等高风险场景中，依赖过时的参数知识可能产生有害输出[7，21]。检索增强生成（RAG）[19]通过在推理时提供外部证据来解决这一差距，其隐含承诺是模型将*忠实地基于*所提供的知识生成输出。当检索到的段落被指定为权威来源（例如，修订后的监管文件或更新的临床指南）时，即使模型自身的参数先验知识存在分歧，模型也应遵循该段落。我们将此属性称为*知识遵从性（KC）*：给定一个被系统视为权威的段落，KC衡量模型遵循该段落的频率，无论事实正确与否——这是一种类似尊重数据库写入的*可控性*指标，不同于在嘈杂或对抗性检索下的准确性。

然而，这一承诺在最需要的时候恰恰瓦解。当检索到的知识与模型的参数记忆相矛盾时——即*知识冲突*场景——LLM表现出不稳定的、不可预测的行为[30，25]。例如，一个基于2024年前医学数据训练的模型可能会忽略检索到的一个描述修订后药物相互作用的段落，反而生成了存储在其参数中的过时推荐。Xu等人[31]表明这并非孤立的失败：模型表现出对参数记忆的*系统性偏好*，并且具有更强参数先验的更大模型对外部证据的*遵从度更低*——这是一种颠覆RAG基本假设的反缩放现象。即使检索是准确的，额外上下文的存在也可能反而增加幻觉倾向，而不是提高忠实度[14]。

我们假设一个关键的促成因素是*架构性*的。在标准RAG中，外部知识和查询token共享同一个自注意力softmax：模型的参数置信度与外部证据在固定的注意力预算上竞争。当冻结的权重在查询token之间产生强烈的键-查询对齐时，对知识token的注意力便被抑制——这是解释为何更大模型更不遵从的一个合理机制（§3.1）。这种机制在*架构上是可以解决的*：通过一个独立的softmax路由外部知识，可以从结构上消除竞争（图1）。

这引出了我们的核心假设：**通过一个不与参数记忆共享注意力预算的独立通道注入外部知识，应该能提高冲突条件下的知识遵从度。** 交叉注意力提供了这样一个通道：它在自身的归一化下计算外部表示上的注意力，从而使自注意力路径不受干扰。我们通过围绕交叉注意力注入构建一个系统，并在五个模型及三个架构系列上进行评估来检验这一假设。

基于这一推理，我们提出了**TokenMem**，一种用于冻结LLM的轻量级记忆系统。**TokenMem**包含三个组件：(1) **TokenMemoryBank**，通过FAISS索引存储和检索知识段落；(2) **交叉注意力适配器**（改编自DecoupledRAG [5]），通过冻结LLM每一层的零初始化门控注入检索到的知识；(3) **两阶段CoT课程SFT协议**，用于训练适配器（约3-7M参数），第一阶段学习知识利用，第二阶段在冲突条件下学习遵从。

高KC是一个可控记忆系统的设计目标：当用户向记忆库写入一个事实时，模型应该使用它，无论其参数编码了什么。

我们的主要贡献是：

- • **系统**：**TokenMem**，一种用于冻结LLM的即插即用记忆系统，通过交叉注意力注入实现高知识遵从。我们在覆盖3个架构系列（Qwen3、LLaMA-3.1、OLMo-3）的5个模型上进行了验证。
- • **训练协议**：我们证明两阶段课程（利用→遵从）是必要的：尽管存在交叉注意力通道，单独的第一阶段会使KC骤降至接近零。这表明该通道为忠实注入提供了*容量*，但需要显式的反事实训练来激活它。
- • **实证观察**：我们观察到标准RAG中知识遵从的反缩放现象——模型越大，遵从度越低，这与我们的架构假设一致——并通过机制分析发现，门控适配器在没有显式监督的情况下学习到了一种冲突感知、逐层特定的注入策略。

## 2 相关工作

#### LLMs中的知识冲突。

当外部上下文与模型的参数知识相矛盾时，LLM表现出不可预测的行为[20]。Xie等人[30]表明模型有时遵循上下文，有时默认使用参数记忆，没有可靠模式；Chen等人[4]证明更丰富的知识源会加剧此类冲突。这些发现促使了两类应对措施：在生成过程中检测和解决冲突的推理时策略，以及通过设计防止冲突的架构方法。**TokenMem**属于后者，提供了一个不与参数表示竞争的独立注入通道。

#### LLMs的知识增强。

多种系统通过外部记忆增强LLM：KBLaM [27]通过矩形注意力注入知识库三元组；MemoryLLM [28]和M+ [29]维护自更新的隐藏状态，但训练所有参数；Knowledge Capsules [15]和FwPKM [32, 17]将知识组织为结构化元组或乘积键条目。这些方法将知识与权重或隐藏状态绑定，使得不重新训练就很难更新。RAG [19]和REALM [9]通过前置段落动态提供知识，但通过相同的自注意力层路由，从而产生冲突。对检索到的段落进行交叉注意力提供了原则上的分离[11, 3]，但需要完整的预训练。DecoupledRAG [5]——最接近的前身——将交叉注意力应用于冻结LLM，但仅在一致知识上评估准确性，而非冲突下的行为。如表1所示，**TokenMem**独特地结合了冻结主干、自由文本输入、交叉注意力注入和冲突意识训练。

表1：知识增强方法的设计空间比较。**冻结**：部署时基础LLM权重不变；**自由文本**：接受非结构化段落（与结构化三元组或潜在状态相对）；**X-attn**：交叉注意力注入路径；**冲突**：包含明确针对知识冲突忠实度的训练或评估。
#### 知识编辑。

ROME [22]和MEMIT [23]直接修改模型权重以更新事实关联，但永久性地改变了模型，并且难以扩展到大型知识库。**TokenMem**则是在推理时从外部库注入知识，允许在不重新训练的情况下进行添加、更新或删除。

## 3 方法

参见标题图1：知识注入路径的比较。**左**：在标准RAG中，知识和查询token共享同一个自注意力softmax，其中参数置信度抑制了知识注意力（AK≤1）。**右**：**TokenMem**通过一个具有自身softmax（∑b=1始终成立）的独立交叉注意力通道路由知识，然后通过一个学习到的门控Wβ进行合并。**TokenMem**通过三个组件（图1）将外部知识注入冻结的LLM：**TokenMemoryBank**（存储和检索知识段落）、**交叉注意力适配器**（通过一个与自注意力分离的通道路由检索到的知识）以及**两阶段课程**（训练适配器即使面对冲突也能忠实地遵循注入的知识）。我们首先形式化标准RAG为何会产生知识冲突（§3.1），然后展示交叉注意力如何避免底层机制（§3.2），最后描述训练过程（§3.3）。

### 3.1 问题：共享路径中的Softmax竞争

考虑一个参数为θ的冻结LLM，它已学会了给定查询x下答案y的参数先验pθ(y|x)。当注入外部知识段落k时，我们关注的是k支持一个与参数先验不同的答案的情况，即知识与模型的“信念”*冲突*。我们将此类样本称为*冲突样本*：形式上，一个满足arg max_y pθ(y|x) ≠ y_k的样本，其中y_k是段落支持的答案。我们将**知识遵从度（KC）**定义为在冲突样本上模型贪婪输出匹配y_k的比例（评估细节见§4.1）。

在标准RAG中，k被拼接进输入：x' = [k; x]。模型通过自注意力处理x'，其中知识token K = {k_1, ..., k_M}和查询token X = {x_1, ..., x_N}共享一个softmax。令k_i表示位置i处的键向量（无论i属于K还是X）。对于查询位置n，分配给知识token j的注意力权重为：¹¹为清晰起见，我们展示完全可见（双向）情况。在因果掩码下，求和限制在≤ n的位置；只要查询侧可见质量为正值，竞争论点仍然成立。

a_{n,j} = exp(q_n^T k_j / √d_h) / (∑_{i∈K} exp(q_n^T k_i / √d_h) + ∑_{i∈X} exp(q_n^T k_i / √d_h)) (1)

对于给定的查询位置n，分配给外部知识的注意力总质量为A_K^{(n)} = S_K / (S_K + S_X) ≤ 1，当S_X > 0时，严格不等式成立。对于固定的知识token得分S_K，A_K^{(n)}随着S_X的增长而*减小*：当冻结权重在非知识token之间产生强烈的键-查询对齐时，共享的softmax会抑制分配给外部知识的注意力比例。这提供了一种机制，可以解释经验观察到的不可靠的知识利用[20, 30]：模型的参数先验与外部证据在固定的注意力预算上竞争。

这种竞争不仅减少了分配给知识的注意力比例，还抑制了其对表示的实际贡献。令知识上下文向量为：

c_n^{rag} = ∑_{j∈K} a_{n,j} v_j (2)

假设值向量有界（即 ||v_j|| ≤ C），根据三角不等式有：

||c_n^{rag}|| = ||∑_{j∈K} a_{n,j} v_j|| ≤ ∑_{j∈K} a_{n,j} ||v_j|| ≤ S_K / (S_K + S_X) · C (3)

在训练良好的大型语言模型中，参数记忆通常会引起强烈的token间关联，导致查询与其自身上下文之间的注意力得分（S_X）支配那些与外部知识token的得分（S_K），即 S_X >> S_K。当内部知识与外部证据冲突时，这种效应尤为显著。在此情况下，界限进一步收紧：

||c_n^{rag}|| ≤ S_K / (S_K + S_X) · C → 0 当 S_X >> S_K (4)

这表明，在参数先验占主导的情况下，标准RAG有效地将外部知识的表示压缩向零，从而