论文精选

PRECOG:SSM状态注入让边缘语言模型实现O(1)上下文检索

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

精选理由

论文把RAG预填充从27秒压到6毫秒,1.2B模型在边缘也能交互式检索,Transformer做不了。

AI 摘要

论文提出PRECOG,把语料库离线预编码为SSM隐藏状态,查询时将最优状态直接注入,省去RAG的上下文重新读取。在带192KB隐藏状态的1.2B参数TENNs-LLM上,PRECOG将prefill延迟从约27秒降到6毫秒以下,约4500倍加速,并保持与上下文RAG相当的答案质量。配套的SMC用认知域聚类构建层级持久记忆,支持保真度/存储调节和O(1)会话初始化。该机制依赖SSM固定长度、与位置无关的循环状态,Transformer的KV缓存因位置纠缠且随上下文线性增长而无法实现。

原文 · arXiv cs.LG

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

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.