TokenMem：面向冻结LLM的忠实知识注入

摘要

检索增强生成（RAG）通过外部知识增强大语言模型（LLM），但存在知识冲突问题：当检索到的信息与参数化记忆相矛盾时，共享的自注意力通路会产生不可预测的输出。我们提出TokenMem，一个轻量级记忆系统，通过专用的交叉注意力通道向冻结LLM注入知识，绕过与残差流中参数化记忆的竞争。TokenMem仅通过两阶段课程训练一个轻薄的门控适配器（约3-7M参数）：先学习一般性知识利用，再在反事实知识下强化忠实遵从。在覆盖三个系列的五个模型（Qwen3-4B/8B/14B、LLaMA-3.1-8B、OLMo-3-7B）的受控实验中，TokenMem在反事实基准上取得69-70%的知识遵从度（Knowledge Compliance，KC），而普通RAG仅为20-52%，差距最高达49个百分点。消融实验表明两阶段课程至关重要：移除第二阶段会使KC骤降至接近零。机制分析揭示，门控适配器在无显式监督的情况下学会了冲突感知、层级特定的注入策略。

TokenMem：面向冻结LLM的忠实知识注入：论文配图 — 图 2：TokenMem 管道概述。 ❶ 记忆库构建与检索：知识段落被编码并存入 FAISS 索引的 TokenMemoryBank；在推理时，查询检索 top-1 段落以生成分层内存表示 𝐦ℓ\mathbf{m}^{\ell}。 ❷ 交叉注意力注入：在每个 Transformer 层，交叉注意力适配器从 𝐦ℓ\mathbf{m}^{\ell} 读取（作为键和值），并通过零初始化的 LoRA 门 WβW_{\beta} 将结果注入到残差流中；仅训练门参数（∼{\sim}3–7M），所有其他组件保持冻结。 ❸ 两阶段课程：第一阶段以知识为基础的 QA 训练，建立交叉注意力路由；第二阶段添加反事实数据，以强制知识冲突下的忠实遵守。