跳到主要内容
4500 倍:当注入状态取代了重新读上下文

4500 倍:当注入状态取代了重新读上下文

图南
图南

· 阅读约 7 分钟

8 月 3 号 arXiv 上挂了篇论文,编号 2608.02560,标题一长串,大意是给边缘端语言模型做结构化记忆。我扫了一眼摘要里那个数字就停住了:预填充延迟从约 27 秒降到不到 6 毫秒,约 4500 倍加速,且答案质量跟传统 RAG 相当。

先别急着激动。这类标题我见过太多了,"4500 倍"往往意味着对照组选得很刁钻——拿最笨的基线跟精心调过的方案比,倍数怎么都好听。但这篇不一样的地方在于,它拆掉的东西不是某个工程环节的优化,而是"上下文必须重新处理一遍"这个默认前提。

说人话就是:你现在用任何一款主流 RAG 或者长上下文模型,查询的时候,系统得把你攒的那堆资料重新塞进模型读一遍——就是预填充(prefill)这道工序。上下文越长,这道工序越贵,Transformer 架构下这个成本随上下文长度线性涨,躲不掉。这是很多人已经默认的物理定律了,像水往低处流一样自然。但这篇论文说:不对,这个前提在状态空间模型(SSM)上根本不成立。

论文提出的机制叫 PRECOG(Pre-Computed Context Injection)。做法一句话就能说清:把文档语料库离线预编码成 SSM 的隐藏状态,查询时直接注入最匹配的那个状态,完事。上下文里不用再摆任何资料原文,也省掉了"重新读一遍"这个动作。

这里有个点非常关键,值得停下来拆一层。你可能会问:这不就是把检索结果提前算好吗?RAG 也可以提前把文档做成 embedding 存起来啊,查询的时候检索 top-k 再塞进上下文,也是"预计算"。

对,但问题就出在"塞进上下文"这一步。RAG 检索回来的文档段落,是要以文本形式摆进 context window,让模型重新读的。这个"重新读"的代价就是 O(L_context) 的预填充计算量——一个字一个字过一遍,跑不掉。而 PRECOG 注入的是隐藏状态,不是文本。SSM 的循环隐藏状态是固定大小的、跟位置无关的,查询时把状态直接注进去,预填充复杂度变成 O(1)。

理解这个区别的关键在于:它注入的不是"资料"本身,而是"模型已经读完了这批资料之后的那种状态"。

这个说法有点绕,我换个方式讲。传统 RAG 好比你去图书馆查资料:你要先在目录里检索,找到几本书,然后坐下了把相关章节一页页翻完,再动笔写答案。PRECOG 是什么,是有人提前把那几本书读完了,写了一份浓缩的读书笔记放在那儿,你查询的时候直接把这本笔记拍在你脑门上,你"相当于"读过那些书了。

这个类比到这儿差不多就该收了——因为我再说下去,你可能会以为 PRECOG 就是做个摘要存起来。不是,它比摘要更接近"状态本身"。摘要还是文本,还是得经过模型重新处理;隐藏状态是模型内部那个循环状态的直接拷贝,是跨过了语言这一层的。这一"跨",就是 O(L_context) 跟 O(1) 的距离。

论文里做了个非常干净的对照:同样的模型(TENNs-LLM,12 亿参数,192KB 隐藏状态),同样的语料,一边用传统上下文内 RAG,一边用 PRECOG 注入状态。答案质量差不多,预填充延迟差了三个数量级。这三个数量级的差,就是"重新读一遍"的代价。

你不能说这是个工程优化。这是架构层面的差异。而且有意思的地方是:这个机制在 Transformer 上不是"不好实现",是"在架构上不可能实现"。KV cache 跟位置纠缠在一起,随上下文长度线性涨——你没法把一个 KV cache 离线算好,然后换个查询直接注进去,因为它的形状取决于当前这条序列本身。而 SSM 的隐藏状态是固定大小的、跟位置无关的,这就给了你"打包带走、随时注入"的可能性。

这是整个论文里最让我觉得反直觉的地方:Transformer 统治了这么多年,我们对"上下文处理成本"的理解全是围绕它建立的——线性增长、缓存命中、窗口管理,全是对着一个位置纠缠的系统在做文章。结果换个架构底座,这些全都不适用了,同一个问题变成了 O(1)。

这里我得插一句自我修正。我一开始看到"pre-fill 从 27 秒降到 6 毫秒",第一反应是"这数字肯定有水",得查了才知道它对比的不只是预填充的总时长,还把整个"把语料塞进上下文"的动作都去掉了。仔细想了想,这么比其实是公允的——因为 RAG 的那 27 秒本来消耗的就是预填充算力。你跑一千个查询,传统 RAG 每一轮都得重新处理一遍同样的上下文(缓存命中可以缓解一部分,但缓存命中本身要求前缀一致,只要有一个文档变了,从变的位置往后的全部重新算);PRECOG 是一次性的离线成本,查询时注入的状态是固定开销。所以 4500 倍不是实验室里挑出来的极端值,它在查询驱动、长期运行的真实场景里只会更明显。

再说 SMC(Structured Memory Consolidation)——论文里跟 PRECOG 搭配的那套分层持久记忆。如果说 PRECOG 解决的是"语料怎么注入"的问题,SMC 解决的是"记忆怎么攒"的问题。它做三件事:把短期情景记忆整合成长期语义记忆、按认知领域聚类、提供一个保真度跟存储的权衡旋钮。这里有个值得展开的细节:SMC 把"记忆"和"语料"放进了同一个空间——不分"这是检索回来的知识文档"和"这是模型记住的对话历史",统一编码成隐藏状态,查询时一起融合。

这个设计细想挺狠的。传统 RAG 有个根上的尴尬:向量检索的"语义距离"和"能不能回答这个问题"是两件事,向量觉得相似的两段话,有时候根本答不上问题(这个类比到这儿已经有点跛脚了,将就看)。而 PRECOG 不走检索-拼接这条路,它把"记忆应该是什么"这个问题从"检索最像的段落"变成了"注入哪个状态最有用"。前者是文本层面的近似匹配,后者是状态层面的直接融合,完全两个粒度。

论文的实验部分不是特别厚实——12 亿参数的模型没法完全代表生产环境里的大家伙,但方向已经够清楚了。而且仔细看作者列表,五位里有几位是做神经形态芯片和边缘硬件的。这其实透露了一点他们为什么要做这件事:边缘端的算力和内存是刚性的,跑不动动辄几百 KB 甚至几 MB 的 KV cache,192KB 固定大小的隐藏状态对它们来说是唯一现实的选择。

这个坑先留着,改天填。PRECOG 还有个显然的软肋:离线预编码意味着知识更新不实时——语料一变,你可能得重新跑一遍离线编码。这个代价怎么跟它的查询时收益权衡,得看具体场景。但那是"工程上还有哪些坑"的问题,不是"这个方向成不成立"的问题。方向是成立的。

说到底这事一句话就能收住:当你不再需要把资料重新读一遍的时候,上下文处理的成本就从"线性读"变成了"直接注入"。懂了这层,你再回头看各家大厂还在卷的"超长上下文窗口"——你会明白他们卷的是怎么让模型读得更多,而这篇论文拆的是"读"这个动作本身。哪边更本质,你自己就能判断。