Structured Memory for Edge Language Models: Persistent Context and Corpus Retrieval via O(1) SSM State Injection

📄 arXiv: 2608.02560v1 📥 PDF

作者: Anusha Madan Gopal, Aras Pirbadian, Kristofor D. Carlson, M Anthony Lewis, Jonathan Tapson

分类: cs.LG, cs.AI, cs.IR

发布日期: 2026-08-03

💡 一句话要点

提出PRECOG以解决检索增强生成模型的预填充成本问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 检索增强生成 状态空间模型 预计算上下文 边缘计算 自然语言处理 实时对话系统 结构化记忆

📋 核心要点

  1. 现有的检索增强生成模型在生成过程中存在较高的预填充成本，影响了实时性。
  2. 本文提出PRECOG机制，通过预计算SSM隐藏状态，降低了查询时的预填充成本。
  3. 实验结果表明，PRECOG在边缘硬件上显著提升了生成速度，预填充延迟减少至6毫秒。

📝 摘要（中文）

检索增强生成（RAG）模型在生成过程中面临着与检索上下文长度成正比的预填充成本，以及在使用Transformer骨干网络时，随着每个生成的token而增长的KV缓存。本文提出PRECOG（预计算上下文注入），通过利用状态空间模型（SSMs）的特性，将预填充成本从$O(L_{context})$降低到$O(1)$。PRECOG在查询时直接注入最佳匹配的SSM隐藏状态，完全绕过了上下文的再摄取。通过在TENNs-LLM模型上进行实验，PRECOG在边缘硬件上将预填充延迟从约27秒减少到小于6毫秒，实现了约4500倍的速度提升。

🔬 方法详解

问题定义：本文旨在解决检索增强生成模型在生成过程中面临的高预填充成本问题。现有方法在生成时需要处理大量上下文，导致延迟显著，尤其是在边缘设备上。

核心思路：论文提出的PRECOG机制通过预先计算文档语料库的SSM隐藏状态，并在查询时直接注入最佳匹配的状态，从而避免了上下文的再摄取，显著降低了预填充成本。

技术框架：整体架构包括预计算阶段和查询阶段。在预计算阶段，文档语料库被编码为固定大小的SSM隐藏状态；在查询阶段，系统根据输入查询直接注入最佳匹配的隐藏状态。

关键创新：PRECOG的核心创新在于利用SSM的固定大小和位置无关的隐藏状态，能够在查询时实现$O(1)$的预填充成本，这在传统的Transformer KV缓存中是无法实现的。

关键设计：在设计中，SSM的隐藏状态被固定为192KB，确保了在查询时的高效性。同时，系统还实现了结构化记忆整合（SMC），结合短期和长期记忆，进一步提升了生成质量。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，PRECOG在TENNs-LLM模型上将预填充延迟从约27秒减少至小于6毫秒，达到了约4500倍的速度提升。这一性能突破使得模型从不可用状态转变为互动式应用，极大地提高了实用性。

🎯 应用场景

该研究的潜在应用领域包括实时对话系统、智能助手和边缘计算设备中的自然语言处理任务。通过显著降低生成延迟，PRECOG可以使得这些应用更加流畅和互动，提升用户体验。未来，该技术还可能扩展到其他需要快速响应的AI系统中。

📄 摘要（原文）

Retrieval-augmented generation (RAG) imposes a prefill cost proportional to retrieved context length, and -- with Transformer backbones -- a KV-cache that grows with each generated token. State-Space Models (SSMs) avoid the second cost by construction; we eliminate the first, collapsing prefill from $O(L_{context})$ to $O(1)$ per query. We introduce PRECOG (Pre-Computed Context Injection), a retrieval mechanism that exploits a property unique to SSMs: the fixed-size, position-agnostic recurrent hidden state is a complete summary of everything the model has read. PRECOG pre-encodes document corpora offline as SSM hidden states and injects the best-matching state directly at query time, bypassing in-context re-ingestion entirely. The same state-injection mechanism enables SMC (Structured Memory Consolidation): a hierarchical persistent memory with cognitive-domain clustering, an adjustable fidelity-vs-storage dial, and $O(1)$ session initialization, which consolidates short-term episodic states into long-term semantic memory and fuses both with retrieved corpus states at query time. We demonstrate the system on TENNs-LLM, a 1.2B-parameter gated-SSM language model with a 192 KB hidden state. PRECOG matches in-context RAG answer quality, reducing prefill latency from $\sim$27 s to $<$6 ms on edge hardware -- a $\sim$4500$\times$ speedup that crosses the threshold from unusable to interactive. The mechanism is architecturally impossible for Transformer KV-caches, which are position-entangled and grow linearly with context length.