ARMT 把「记忆」焊进 Transformer:用恒定显存换无限上下文
长上下文是这两年 LLM 工程化最大的「显形成本」。标准 Transformer 注意力是 O(n²) 计算、O(n) 显存——上下文从 8K 推到 128K,推理成本要涨两个数量级。现有解法几乎都集中在 KV Cache 压缩(JLT、MosaicKV)或混合注意力(HydraHead、DeepSeek V4),都在「保留一段连续 KV」假设里打转。
7 月 13 日挂在 arXiv 的 2607.11614(Kuzmin 等 11 位作者,MIPT + AIRI + MBZUAI 等机构合作)提出不太一样的角度:Associative Recurrent Memory Transformer (ARMT)——把联想式循环记忆直接插入 Transformer 某些层,实现恒定显存扩展上下文。
核心贡献
- 两个领域特定长上下文数据集,为窄域微调场景设计,比通用 needle-in-haystack 更接近真实工作负载;
- 完整训练配方:持续预训练 + 合成长文数据 + 课程学习 + 选择性把联想记忆接入部分层(不是全量替换);
- 系统实验:ARMT 增强后能处理远超原生窗口的输入且不退化,对分布外长度泛化更好,且在原始窗口内少用 30% FLOPs。
个人评论
ARMT 不是范式颠覆,更像「在 Transformer 内部挖一个 LRU 缓存」。精神上接近 SSM/Mamba,但 ARMT 显式选择「保留原主干 + 局部加记忆」的渐进路线,工程门槛低很多。
30% FLOPs 降幅对已上线长上下文应用是直接成本收益;「选择性插入层」是关键工程洞察——全量替换会破坏表征局部性,稀疏插入在效率与质量间找到 trade-off;论文强调「窄域微调场景」,说明 ARMT 不是通用 1M 上下文基模,而是面向垂直长文工作负载(法律、医疗、代码库)的实用路径。
接下来要看:ARMT 的「联想记忆」是否会和 KV Cache 压缩方法叠加——「外存 + 内压」组合可能是长上下文成本下降的下一站。
从 ReContext 到 JLT 再到 ARMT,长上下文已经不是「能不能做」的问题,而是「用多少成本做」的问题。每篇都在把单位 token 推理开销往下压一档——这才是 LLM 真正走进生产环境的底层燃料。