让大模型高效处理超长文本，不重训练也不占显存
MoNe: Modular Neural Memory for Efficient Long Context Inference
- 用模块化记忆网络分段读取上下文，推理时只用查询生成记忆
- 128K token时节省80%算力和显存，仅增6.4%参数量
- 适合需要长文本理解的场景，如信息检索、文档分析
我们提出 MoNe，一种轻量级模块化神经记忆机制，可附加到任意冻结的预训练 Transformer 上，实现无需重训练的长上下文推理。MoNe 通过测试时学习的快速权重神经记忆网络，以固定大小分段读取上下文，并采用层局部梯度更新；推理时，仅凭查询令牌生成记忆键值对，无需重新读取上下文。该两阶段设计使推理开销与上下文长度解耦，在 128K token 时实现约 $O(N)$ 的预处理和 $O(1)$ 的查询成本，峰值 GPU 显存不随 $N$ 增长。相比 ICL，MoNe 在 128K token 下降低约 80% 的计算量与峰值显存，仅增加 6.4% 参数。其性能在 RULER 数据集的 needle-in-a-haystack 与词提取任务中显著优于 ICL，且能推广至远超主干模型原生窗口长度的上下文。
原文摘要 · Abstract (English)
We present MoNe, a lightweight modular neural memory that attaches to any frozen pretrained Transformer to enable long-context inference without retraining. MoNe reads context in fixed-size segments via test-time learning of fast-weight neural memory networks with layer-localized gradient updates; at inference, the memory generates keys and values from the query tokens alone, with no context tokens re-read. This two-phase design decouples inference cost from context length, achieving $O(N)$ preprocessing and $O(1)$ query cost with peak GPU memory that does not grow with $N$. At 128K tokens, MoNe reduces both compute and peak GPU memory by approximately 80% compared to ICL with only 6.4% parameter overhead. MoNe generalizes to context lengths far beyond the backbone's native window, achieving strong performance on needle-in-a-haystack and word extraction benchmarks from RULER, where ICL degrades sharply.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。