高通提出外挂式在线可写神经记忆模块，长上下文显存占用大降 80%！

智猩猩AI整理
编辑：BugMaker

长上下文越做越长，一个最直接的问题也越来越明显：

模型到底有没有必要在每次提问时，都重新把几十万Token从头读一遍？

传统In-Context Learning的做法很直接，把完整上下文和问题一起塞进Transformer。这样虽然保留了全部信息，但问题Token需要直接和大量上下文Token做Attention，随着上下文不断变长，计算量和KV Cache占用也跟着上涨。

RAG走了另一条路，先检索与问题最相关的少量文本，再交给模型回答。

但如果答案依赖分散在多个位置的信息，仅靠Embedding相似度检索Top-K文本，也可能遗漏关键证据。

来自Qualcomm AI Research的研究团队最近提出了MoNe（Modular Neural Memory），尝试给已有Transformer外挂一个“可在线写入的神经记忆”。该工作发表于第43届国际机器学习大会（ICML 2026）的AdaptFM Workshop，并收录于PMLR 306。

它最关键的改变是：

长上下文只需要分段读一遍，把信息写入可更新的神经Memory。真正回答问题时，不再重新读取原始上下文。

论文使用冻结的Qwen2.5-0.5B-Instruct作为Backbone。MoNe只增加约6.4%的参数开销，离线训练阶段最长只见过4K上下文，却能够直接测试到128K。

在128K上下文下，相比ICL，MoNe将峰值GPU显存和总FLOPs都降低了约80%。

更重要的是，128K下MoNe在S-NIAH、MK-NIAH和Frequent Word Extraction三个RULER任务上的对应指标分别达到0.96、0.94和0.96。

01

128K上下文不只更贵，

ICL还可能越读越差

长上下文模型通常有两个问题。

第一个是成本。

对于标准Self-Attention，Context越来越长，Attention计算量会以O(N²)快速增长。与此同时，KV Cache占用也会随着上下文长度继续增加。

论文给出的结果非常直观。

32K上下文下，ICL需要2.48 GB峰值GPU显存和58.06 T FLOPs。

到了128K，峰值GPU显存进一步增加到7.07 GB，计算成本也随上下文长度快速上升。

但比成本更关键的是另一个问题：

上下文放得进去，不代表模型真的能够有效使用。

在S-NIAH任务中，ICL在32K时还能达到0.94，48K下降到0.64，128K只剩0.28。

MK-NIAH更明显，从32K的0.93一路下降，到了128K直接变成0.00。

这正是MoNe想解决的问题：

能不能既不让模型每次重新面对几十万Token，又尽量保留分散在长上下文中的信息？

RAG可以减少真正送入模型的Token，但面对分散在长Context中的信息，Chunk检索仍可能遗漏必要证据。

MoNe因此换了一条路线：

不再只从长Context里挑出一部分，而是尝试把上下文逐步写入神经记忆。

02

不改Transformer骨干，MoNe给

每一层外挂一块神经记忆

MoNe首先有一个重要的设计前提：

原来的Transformer Backbone始终保持冻结。

论文实验使用的是Qwen2.5-0.5B-Instruct。MoNe没有重新训练一个新的长上下文Transformer，而是在每个Decoder Layer旁边增加一个Fast-Weight Neural Memory。

这里的Fast Weight可以简单理解成：

这是一组允许在Test-Time Learning阶段，根据当前上下文快速更新的参数。

论文中，每一层的Memory采用一个SwiGLU MLP。

整个长Context首先被切成固定大小的Segment，默认每段为512 Token。

以128K上下文为例，它不会作为一个完整的128K序列一次性送入Attention，而是按照512 Token连续切分并逐段处理：

Segment 1 → 更新Memory → Segment 2 → 更新Memory → …… → 得到最终Fast-Weight State

核心在于，每一层Memory的更新都是局部完成的。

MoNe利用当前Transformer Layer产生的Key和Value，通过Associative Memory Loss把对应关系写入这一层自己的Fast Weight。简单来说，就是让Memory在给定Key时，能够生成尽可能匹配对应Value的输出。

更关键的是，某一层更新Memory时，只依赖这一层自身Forward产生的信息，不需要把梯度穿过后续Transformer Layer。

因此，MoNe可以模块化地挂接在已有Transformer上，而不用为了每次Memory更新重新反向传播整个Backbone。

不过，这里的“无需重新训练”需要准确理解。

MoNe不需要重新训练或更新原有Transformer Backbone的权重，但MoNe自身并不是完全零训练。

LoRA Adapter、学习率相关参数、Momentum Projection等Meta Parameter仍需要提前离线训练。部署后这些参数保持冻结，Test-Time Learning阶段只继续更新Fast Weight。

03

上下文只写一次，后续Query
不再重读128K Token

MoNe和传统ICL真正拉开差别的地方，发生在Context全部处理完以后。

传统ICL到了提问阶段，问题Token仍然需要直接面对完整Context。

MoNe不是。

Context已经在前面的Test-Time Learning阶段，被逐段写进各层Fast Weight。

因此真正出现Query时，每一层都利用Query从Fast-Weight Memory中生成对应的Memory Token，再把这些Memory Token投影成额外的Key和Value加入Attention。

换句话说：

Query不再直接访问之前的128K原始Context，而是通过神经记忆生成的Memory Token读取已经写入的信息。

这让MoNe形成一个明显的两阶段结构：

第一阶段是写Memory。

长度为N的上下文被固定大小Segment逐步处理，因此总预处理成本随N线性增长，可以看成O(N)。

第二阶段才是读Memory回答问题。

这时候原始Context已经不参与Attention，因此相对于Context长度N，Query阶段的成本不再随N增长，可以看成O(1)。

同时，每层Memory KV始终保持固定的T个Entry，不会随着总Context长度继续增长。

这也是为什么论文从32K扩大到128K以后，MoNe的Test-Time Learning峰值显存仍然保持在1.41 GB。

Inference阶段的峰值GPU显存则为1.29 GB。

另一个实际优势是：

同一份Context写入完成以后，最终Fast-Weight State还可以被后续多个Query重复使用。

如果又出现新的Context，也可以在已有Memory基础上继续更新，而不用重新处理之前的内容。

04

只训练到4K却外推128K，
MoNe在三项RULER任务仍保持高分

论文最终在RULER的三个任务上进行了测试：

S-NIAH负责单个Key-Value检索。

MK-NIAH在不同位置放入多个Key-Value，需要抵抗干扰信息。

Frequent Word Extraction则要求从整个上下文中找出出现最频繁的三个目标词。

其中S-NIAH和MK-NIAH采用Sub-EM作为指标，Frequent Word Extraction采用Variable Recall，均为越高越好。

Baseline包括完整上下文ICL，以及基于BGE-Large检索的RAG。后者将Context切成128 Token的Chunk，并测试K=1、4、8三种检索数量，在每个Context Length下报告最佳结果。

真正拉开差距的，是Context超过模型原生32K窗口以后。

ICL的表现随着上下文继续变长快速下降，RAG相对稳定，但在需要更全面利用长上下文信息的任务上仍与MoNe存在明显差距。

MoNe则一直测试到了128K，并在三项任务上都保持了较高水平。

更关键的是，MoNe离线训练时最长只见过4K上下文，128K已经相当于把测试长度扩展到了训练长度的32倍。

论文将这种长度泛化部分归因于Segment-local RoPE。简单来说，每个512 Token Segment内部都重新使用一套固定的位置范围，因此Context继续变长时，不需要把位置编码一起拉长。

效率上的差距同样会随着Context变长被放大。

到了128K，MoNe的峰值GPU显存和总计算量都只剩ICL的大约五分之一。

而且MoNe的峰值显存并没有随着Context从32K扩大到128K继续增长，这与固定大小的Memory以及固定长度Segment逐段处理有关。

作者还测试了一个问题：Memory到底需要挂多少层？

结果很明显。只给模型最后几层增加Memory，短上下文影响不大，但Context一旦拉长，性能会迅速下降。覆盖全部24层时，128K下的效果最好。

这组实验说明，在MK-NIAH上，Memory覆盖更多Decoder Layer，更有利于维持长上下文下的性能。

Segment Size也有类似现象。

每次切得太短，虽然能少算一点，但Context越长，记忆效果下降得越明显。相比之下，论文默认采用的512 Token只增加了有限的计算成本，却能明显改善128K下的表现。

因此论文最终选择了全部24层 + 512 Token Segment作为默认配置。

MoNe真正有意思的地方，不只是把长上下文再次做了一次压缩，而是尝试把原本保存在Prompt里的历史信息，转化为可以持续更新并重复读取的神经记忆状态。

这意味着长上下文推理可以被拆成两件事：第一次如何高效写入Memory，以及后续Query如何不再反复读取原始Context。

不过目前实验仍集中在Qwen2.5-0.5B和RULER受控任务。论文也明确将更大模型、真实单/多文档QA和长期对话历史留作后续验证，因此现在还不能直接把128K上的结果外推到真实复杂长文档场景。

MoNe更像是在验证一条路线：长上下文不一定只能继续扩大Context Window，也可以被逐步写入一块固定大小、可更新的神经记忆。

关注+星标，获取AI前沿进展与开源一线动态