用状态空间模型实现零开销检索,让边缘设备语言模型瞬间响应。
Structured Memory for Edge Language Models: Persistent Context and Corpus Retrieval via O(1) SSM State Injection

- 将文档库预编码为固定大小的状态,查询时直接注入,无需重新读取。
- 在边缘设备上将预填充延迟从27秒降至6毫秒,提速4500倍。
- 支持长期记忆存储与快速会话初始化,适合资源受限的实时应用。
检索增强生成(RAG)在生成前需处理全部检索内容,导致预填充成本随上下文长度线性增长;而基于Transformer的模型还需维护随生成过程不断增长的键值缓存。状态空间模型(SSMs)天生避免第二类开销;本文提出PRECOG(预计算上下文注入),通过利用SSM独有的特性——固定大小、位置无关的循环隐藏状态可完整总结已读内容,将预填充复杂度从O(L_context)降至O(1)。PRECOG离线将文档库编码为SSM隐藏状态,并在查询时直接注入最佳匹配状态,完全跳过上下文重读。该机制还支持SMC(结构化记忆整合):一种分层持久记忆系统,具备认知域聚类、可调保真度-存储比、以及O(1)会话初始化能力,可将短期事件状态融合为长期语义记忆,并在查询时与检索到的文档状态结合。实验在1.2B参数的门控SSM模型TENNs-LLM上进行,其隐藏状态大小为192 KB。PRECOG达到与上下文内RAG相当的答案质量,将边缘设备上的预填充延迟从约27秒降至不到6毫秒,提速约4500倍,使系统从不可用变为可交互。该机制对位置耦合且线性增长的Transformer KV缓存而言在架构上不可能实现。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。