🤖 machine learning

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

本文介绍了 PRECOG 和 SMC，这是一种针对状态空间模型（State-Space Models）的检索增强生成系统，通过将文档语料库预编码为固定大小的隐状态来实现 $O(1)$ 的上下文注入，从而使交互式边缘语言模型能够实现相比传统基于 Transformer 的 RAG 的大幅加速。

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

发布于 2026-08-04

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

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

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图教一个超级聪明的机器人通过喂给它海量的书籍来回答问题。在人工智能的世界里，最流行的方法就像是递给机器人一叠纸并对它说：“在回答我的问题之前，把这些纸上的每一个字都读完。”这种方法虽然可行，但速度很慢。如果这叠纸非常厚，机器人为了读完这些上下文，甚至在开口说“你好”之前就要花很长时间。此外，机器人必须记住它刚刚读过的每一个字，这非常占用内存，就像试图在说话时把整部百科全书都记在脑子里一样。

科学家们一直在寻找一种更快速、更轻量化的方法，特别是针对那些运行在手机或智能家电等小型设备上的机器人，在这些设备上，内存有限且速度至关重要。他们正在探索一种不同类型的脑架构，称为“状态空间模型”（State-Space Model，简称 SSM）。不要把 SSM 想象成一个逐页阅读书籍并背诵每个单词的机器人，而要把它想象成一个阅读书籍后，能立即将整个故事压缩成一张微小的、神奇的“总结卡”的机器人。这张卡片承载了所读内容的“意义”，但并不包含单词本身，而且它并不在意故事发生在书中的哪个位置。研究人员提出的核心问题是：我们能否利用这个“总结卡”的小技巧来跳过缓慢的阅读过程？如果我们能做到，我们就能制造出能够瞬间回答问题的 AI，即使是在微小的设备上，也不需要依赖云端庞大的计算机。

这篇论文介绍了一个巧妙的新技巧，叫做 PRECOG（意为预计算上下文注入），它证明了“是的，我们可以”。作者们通过一种特定类型的 AI 模型——TENNs-LLM，找到了如何将一段文本通过 AI 运行一次，从而生成那张微小的“总结卡”（隐藏状态）并将其保存下来。随后，当用户提出问题时，系统不再让 AI 重新阅读文本，而是直接取出保存好的总结卡，并将其直接“丢进”AI 的大脑中。这就像是跳过了考试中的阅读环节，直接把学生之前复习好的笔记递给了他。

实验结果在速度方面令人惊叹。在团队测试的特定边缘侧硬件上，旧的“阅读”上下文的方法需要大约 27 秒才能得到第一个答案。而使用 PRECOG 后，这一时间降到了不到 6 毫秒。这实现了约 4,500 倍的加速，将一个原本无法使用的过程变成了即时且具有交互性的体验。论文从数学上证明了这种捷径并不会让 AI 变“笨”；其答案与逐字阅读文本的效果一样好，因为这种特定类型的 AI 工作方式允许总结卡成为文本的完美替代品。

论文还介绍了第二个功能，称为 SMC（结构化记忆整合）。如果说 PRECOG 是关于记住一座图书馆的书籍，那么 SMC 就是关于记住一个人的生活。它将 AI 的长期记忆（例如你的偏好或过去的对话）组织成不同的“文件夹”或簇，比如“情感”、“事实”或“地点”。这样，当你开始一段对话时，它就能立即将正确的文件夹加载到 AI 的大脑中，让 AI 记住你是谁以及你喜欢什么，而无需你每次都提醒它。整个系统旨在小型、低功耗设备上运行，使得在你的口袋里拥有一个超智能、记忆丰富的 AI 助手成为可能。

作者对他们的数学推导非常有信心。他们证明了对于这种特定类型的 AI，这种“总结卡”方法并非一种近似或猜测，在数学上与阅读文本是完全等同的。然而，他们也指出，这个技巧仅适用于这种特定类型的 AI 架构。更常见的 AI 模型（称为 Transformer）构建方式不同；它们会纠缠于单词的顺序，因此你不能直接用总结卡来替换文本。对于那些模型，你仍然必须每次都阅读文本。但对于未来的小型化、快速化、基于边缘侧的 AI 而言，这篇论文展示了一条实现强大记忆力而无需承担沉重阅读负担的路径。

技术摘要：通过 O(1) SSM 状态注入实现边缘语言模型的结构化记忆

1. 问题陈述

在边缘设备上部署检索增强生成（RAG）目前受到主流 Transformer 架构两个结构性限制的约束：

线性预填充成本（Linear Prefill Cost）： 上下文内 RAG 需要在预填充阶段逐个 token 地重新摄取检索到的上下文 token。这造成了与上下文长度（ $O(L_{context})$ ）成正比的延迟和能量成本，这对于带宽受限的边缘硬件来说是难以承受的。
位置纠缠的 KV 缓存（Position-Entangled KV Caches）： Transformer 的键值（KV）缓存包含了位置编码（旋转或绝对位置），这使得它们无法在不同上下文或位置之间复用。此外，缓存大小随上下文长度线性增长，造成了随语料库规模扩展性极差的存储瓶颈。

状态空间模型（SSM）提供了一个理论上的替代方案，通过将先前的上下文压缩为固定大小、与位置无关的循环隐藏状态。然而，现有的 RAG 实现仍依赖于逐 token 摄取，未能利用 SSM 独特的代数特性来完全消除预填充成本。

2. 方法论

本文引入了 PRECOG（预计算上下文注入） 和 SMC（结构化记忆整合），这两个机制利用了 SSM 的时间平移不变性（time-translation invariance）来绕过上下文摄取过程。

PRECOG：预计算上下文注入

PRECOG 通过离线预计算检索文档的隐藏状态并在查询时直接注入，从而重新定义了 RAG。

离线索引： 文档语料库被划分为若干块（chunks）。SSM 语言模型处理每个块以捕获每一层的最终隐藏状态（ $\bar{h}^{(i)}$ ）。这些固定大小的状态（对于 1.2B 参数模型为 192 KB）存储在闪存中，并配有轻量级的句子编码器键（keys）用于检索。
查询时注入： 在接收到查询后，系统通过其键检索匹配度最高的 top- $k$ 个块。系统不再重新摄取文本，而是加载相应的预计算隐藏状态，并将其作为模型的初始循环状态（ $\bar{h}_{init}$ ) 进行注入。随后，模型仅处理查询部分的 tokens。
理论保证（定理 1）： 本文证明了对于任何初始状态 $h_0$ 、上下文 $c$ 和查询 $q$ ，状态轨迹 $S(h_0, c \oplus q)$ 在代数上等同于 $S(S(h_0, c), q)$ 。因为 SSM 的更新仅取决于当前状态和 token（而非绝对位置），注入预计算状态在数学上等同于顺序处理了上下文 tokens。这保证了 PRECOG 能精确匹配上下文内 RAG 的回答质量（仅存在量化噪声）。

SMC：结构化记忆整合

SMC 将状态注入机制扩展到持久化的设备端记忆，用于累积用户偏好、交互历史和日志。

层次化聚类： 交互块被路由到认知领域簇（例如：情感、时间、空间、事实）中，并根据召回模式（情节式 vs. 语义式）进一步细分为子簇。
保真度-存储调节旋钮（Fidelity-Storage Dial）： 参数 $K$ $K$ 控制每个块保留的状态数量：
- $K=N_c$ ：无损情节记忆（保留所有状态）。
- $K=1$ ：语义记忆（仅保留最终状态，与 PRECOG 块一致）。
O(1) 会话初始化： 在会话开始时，系统识别主导领域并直接将相应的整合语义状态注入 SSM。这使得模型在无需重新摄取任何 token 的情况下，即可具备历史交互的初始化状态，无论累积的历史长度如何，均保持 $O(1)$ 延迟。

3. 核心贡献

RAG 的算法重构： 本文证明了通过注入预计算的 SSM 隐藏状态，可以将上下文摄取成本从 $O(L_{context})$ 降低到 $O(1)$ ，这一机制由于位置纠缠问题在 Transformer 架构中是结构性不可能实现的。
等价性证明： 定理 1 提供了形式化保证，证明 PRECOG 产生的状态轨迹与上下文内 RAG 完全一致，确保了加速过程不会以牺牲回答质量为代价。
统一的记忆基质： 同一种状态注入原语统一了静态语料检索（PRECOG）和动态持久化记忆（SMC），允许边缘模型在单个查询步骤中融合检索到的知识和用户历史。
边缘部署可行性： 通过在 TENNs-LLM（一个具有 192 KB 总隐藏状态的 1.2B 参数门控 SSM）上实例化这些方法，作者展示了状态级检索在神经形态边缘硬件上的实用性，其微小的状态足迹符合片上内存的约束。

4. 结果

系统在 SQuAD v1.1 开发集上使用 1.2B 参数的 TENNs-LLM 骨干网络进行了评估。

质量： PRECOG（top-1 注入）实现了 58.0 的精确匹配（EM）和 73.4 的 F1 分数，在 FP16 量化噪声范围内与上下文内 RAG 基准（EM 58.2, F1 73.6）相匹配。这从经验上验证了理论上的等价性。
延迟： 在边缘硬件（模拟神经形态处理器）上，PRECOG 将预填充延迟从 ~27 秒（摄取 512 个 token 块所需的时长）降低至检索和注入开销的 <6 毫秒。包括摄取查询 tokens 在内的总首 token 时间（TTFT）约为 585 毫秒。文中提到的 ~4,500 倍加速 特指消除了约 27 秒的上下文摄取阶段，而非总查询延迟。
存储效率： 对于 512 个 token 的块，PRECOG 状态需要 192 KB，而同等的 Transformer KV 缓存则需要 ~16 MB（减少了 85 倍）。对于 10,000 个块的语料库，PRECOG 需要 1.9 GB 存储空间，而假设的 KV 缓存方法则需要 160 GB，后者在消费级设备上是不可行的。
消融实验：
- 注入深度： 仅注入底层 12 层（96 KB）即可保留 99% 的全层 F1 分数，表明存在进一步优化存储的空间。
- 块长度： 对于长度在 ~600 个 token（模型的有效记忆视界）以内的块，性能与上下文内 RAG 一致。超过此长度后，随着 SSM 自然地遗忘远端 token，PRECOG 的表现会平滑下降，这与具有相同骨干网络的上下文内 RAG 行为一致。

5. 意义与主张

论文声称 PRECOG 和 SMC 代表了边缘语言模型处理检索与记忆方式的根本转变。其重要性在于：

架构排他性： 该机制依赖于 SSM 的时间平移不变性，这使得它对于 Transformer KV 缓存而言是“架构上不可能实现的”。
边缘可行性： 它解决了目前阻碍交互式 RAG 在带宽受限设备上运行的“预填充瓶颈”，实现了从依赖云端的推理向设备端、低延迟交互的转变。
设计原则： 作者认为，随着 SSM 缩小与 Transformer 的质量差距，其结构特性（固定大小、与位置无关的状态）将催生出一类全新的检索、缓存和记忆扩展算法，而这些算法是基于注意力机制的架构所无法触及的。

论文对局限性保持了谦逊的态度，承认 PRECOG 继承了 SSM 有限的记忆视界（无法恢复骨干网络会遗忘的信息），并且 top- $k$ 组合启发式算法（用于多块检索）是缺乏精确性保证的近似方法。此外，神经形态部署的数据基于 FPGA 测量和 12 nm 电力模拟，尚需等待物理流片验证。