Knowledge Graph Enhanced Memory-Augmented Retrieval for Long Context Modeling

技术摘要：KGERMAR

问题陈述
长文本语言建模面临着一个根本性的挑战，这不仅仅是单纯地扩展上下文窗口，而是如何在数千个 token 中维持对实体状态和关系的连贯理解。标准方法存在两个主要局限性。首先，扩展注意力模型（如 LongLoRA、YARN）存在“迷失在中部”（lost-in-the-middle）的退化现象，即位于长上下文中间的信息会被系统性地降低关注度。其次，记忆增强检索系统（如 MemLong、ERMAR）完全依赖于语义相似性。虽然它们在检索具有匹配术语的段落方面非常有效，但无法区分那些仅仅是偶然提到相关术语的段落与那些描述了相同实体之间明确的因果或结构关系的段落。这种局限性在需要结构化推理的场景中至关重要，例如在长达整个故障排除过程的设备错误状态演进追踪中，相关的上下文可能在词汇上与查询没有任何重叠，但在深层结构上具有关联。

方法论
本文提出了 KGERMAR（知识图谱增强的记忆增强检索），这是一个通过在推理过程中构建动态的、特定于上下文的知识图谱，以整合图结构嵌入与文本语义的框架。该系统通过四个集成组件运行：

1. 上下文知识图谱构建：

   • 提取： 系统使用基于 BERT-large（在 CoNLL-2003 上微调）的实时命名实体识别（NER）和基于 BERT-base（在 TACRED 上微调）的关系提取。
   • 图构建： 它在滑动窗口内提取实体和关系（例如 CAUSEDBY、PARTOF、USEDFOR）。为确保质量，它采用了实体合并（通过字符串和嵌入相似性合并变体）、关系置信度评分（结合分类器确定性和频率）以及图过滤。
   • 动态特性： 与静态方法（如 ConceptNet）不同，该图是从输入文本本身构建的，能够捕捉领域特定和会话特定的实体（例如，特定的打印机型号及其唯一的错误代码），而这些实体在预构建的知识库中并不存在。

2. 图增强嵌入生成：

   • 编码： 提取的实体使用关系图卷积网络（R-GCN）进行编码，通过多跳消息传递来捕捉个体属性和结构角色。
   • 融合： 这些图结构嵌入通过跨模态注意力机制与文本实体嵌入（源自基础语言模型）进行融合。这创建了一种混合表示，既捕捉了表层语义，也捕捉了深层的关系模式。

3. 多组件记忆架构：

   • KGERMAR 维护三个专门的记忆库：
     • 上下文记忆（Contextual Memory）： 存储来自语言模型层的 token 级键值对（遵循 ERMAR）。
     • 语义记忆（Semantic Memory）： 存储非重叠块的稠密文本嵌入（使用 BGE-M3）。
     • 结构记忆（Structural Memory）： 存储融合后的文本-图实体嵌入。
   • 所有记忆库都通过增量方式更新，并在达到容量时通过 LRU 驱逐策略进行管理。

4. 混合检索机制：

   • 对于给定的查询，系统在所有三个记忆库中执行并行 Top-K 检索。
   • 检索信号使用学习到的融合权重 ($\lambda_1, \lambda_2, \lambda_3$) 进行组合，以产生最终得分。
   • 检索到的表示通过检索因果注意力（retrieval causal attention）注入到指定高层的语言模型中。

训练目标
该模型采用多任务学习来同时优化四个目标：

• 语言建模损失 ($L_{LM}$)： 使用检索到的上下文进行下一 token 预测。
• 知识图谱嵌入损失 ($L_{KG}$)： 通过二元交叉熵进行链路预测，以训练 R-GCN。
• 检索排序损失 ($L_{retrieval}$)： 使用基于间隔的排序（margin-based ranking），以确保相关的上下文比不相关的上下文得分更高。
• 跨模态对齐损失 ($L_{alignment}$)： 确保文本表示与结构化表示之间的一致性。

核心贡献

1. 动态知识图谱构建： 一个能够在推理过程中从输入文本实时提取领域特定实体和关系的流水线，实现了对不存在于静态知识库中的专业内容的适应。
2. 混合多组件记忆： 一种整合了上下文、语义和结构化表示的架构。其结构记忆库专门捕捉了与纯语义相似性互补且通常更具判别力的实体关系。
3. 实证验证： 证明了引入结构化知识能显著提升长上下文建模能力，与强基准模型相比，困惑度（perplexity）降低了高达 8.5%，且内存效率提高了 2–2.5 倍。

实验结果
该框架在 SlimPajama（84.7K 样本用于微调）和三个基准数据集 WikiText-103、PG-19 及 Proof-pile 上进行了评估。

• 困惑度（Perplexity）： 在 1K 到 32K token 的上下文长度范围内，KGERMAR 在所有记忆增强型 3B 模型中实现了最低的困惑度。在 WikiText-103 上，它显著优于 ERMAR 和 MemLong（例如，在 16K token 时，困惑度为 7.14 vs 7.80）。
• 上下文学习（ICL）： 在五个 NLU 任务（SST-2, MR, Subj, SST-5, MPQA）上，KGERMAR 在 4-shot（78.6%）和 20-shot（81.2%）设置下均取得了最高平均准确率，超越了 ERMAR 和 MemLong。在需要实体级辨析的任务（如 Subj, MPQA）中，增益尤为显著。
• 内存效率： KGERMAR 展示了卓越的内存效率，其内存使用量比基准模型低 2–2.5 倍。在 32K token 时，其峰值内存比 ERMAR 少 36%，比 MemLong 少 41%。随着上下文长度的增加，单位 token 的归一化内存使用量有所改善，验证了其缓解标准注意力机制中二次方内存增长的能力。
• 缩放行为： 虽然由于实时图构建的存在，KGERMAR 的单 token 延迟略高，但它表现出随序列长度近乎线性的延迟增长，这使得它在长上下文部署中比 ERMAR 更具可预测性。

意义
本文声称，KGERMAR 通过将依赖从固定的静态知识库转向领域自适应知识表示，推动了记忆增强语言建模的发展。通过从输入上下文中动态构建知识图谱，该系统能够对仅靠语义相似性无法捕捉的实体状态和因果演进进行推理。这种方法为在通用硬件上进行长上下文建模提供了一个切实可行的解决方案，在提供高质量建模的同时，显著降低了内存开销，这对于需要结构化推理的技术支持、科学文献和法律文档分析等应用至关重要。