历史重现:释放线性注意力的长上下文能力 History Replays Itself: Unleashing Long-Context Capability in Linear Attention
免训练稀疏记忆让线性注意力百万 token 检索从 1.7% 升至 70.3%
前置知识
线性注意力与固定尺寸循环状态
线性注意力类模型(Gated DeltaNet、Mamba-2 等)不像 softmax 注意力那样为每个历史 token 保留键值对,而是把已读历史压缩进一个固定尺寸矩阵状态 $S_t$,按 $S_t = S_{t-1}A_t + \beta_t v_t k_t^{\top}$ 更新,用 $o_t = S_t q_t$ 读取。状态尺寸与序列长度无关,因此计算与显存随长度只线性增长。代价是:历史只能以有损压缩的形式留在 $S_t$ 中。
论文的全部出发点就是这个固定状态被迫同时承担持续计算与长期记忆两个角色,导致百万 token 检索崩溃;理解状态更新机制才能明白 NoteBank 为何把证据存档移出循环状态。
MHA 的 KV 缓存
标准 softmax Transformer 推理时为每个历史 token 缓存键向量 $k$ 和值向量 $v$,缓存字节数约为 $4|L|DP$($|L|$ 层数、$D$ 隐藏宽度、$P$ 序列长度、BF16 两字节)。在与本文 GDN 同宽度、同深度、同精度、1M token 的设置下,全注意力 KV 缓存约需 192 GiB 显存。
KV 缓存是论文存储效率对比的基准:NoteBank 的持久上下文存储仅约 74.5 MiB,即其 1/2600 左右,这是『近乎免费午餐』论断的核心数据。
Delta 规则与 Gated DeltaNet
Delta 规则把状态当关联记忆:写入新键值对前先减去同键的旧内容再秩一更新,减少记忆干扰。Gated DeltaNet 在此之上引入门控衰减矩阵 $A_t$ 和写强度 $\beta_t$,精细控制保留与覆写。本文的 GDN 主干为 1.3B 参数、24 层、每层 8 头,共 192 个被索引的层-头通道。
头条实验(RULER@1M 1.7%→70.3%)就在冻结 GDN 上完成;且 NoteBank 的写入与寻址直接复用其原生的 key/query 特征,必须了解这一架构。
Mamba-2 与长上下文状态爆炸
Mamba-2 是选择性状态空间模型,用输入依赖的线性门控更新状态。论文发现在数十万 token 以上的推理中其循环状态幅度会爆炸、读出失效,需要在回答前用 $S \leftarrow \lambda S$ 重缩放状态(实验取 $\lambda=0$,即直接清零并保留卷积缓存与位置信息)才能可靠作答。
这解释了 Mamba-2-8B 在 MemoryAgentBench 上 0.0 分的原因,也展示了解耦记忆的独特红利:状态清零后证据仍可从外部存档重放回来。
RULER 与大海捞针(NIAH)
RULER 是长上下文评测基准,其中 NIAH 任务把关键信息(针)埋入长干扰文本(草堆)并要求精确取出。本文使用 6 个任务:单针 S-NIAH-1/2/3 与 multikey、multiquery、multivalue,每任务 100 个固定样本,文档长度从 4K 到 1M token。
1.7%→70.3%、+68.6/+66.6/+37.8 个百分点等所有头条数字都来自这 6 个任务的均值,是衡量『循环模型能否恢复远程检索』的主战场。
LoRA 适配器与读阶段适配
LoRA 冻结原权重,仅在选定投影上添加低秩更新 $\Delta W = (\alpha/r)BA$,$B$ 零初始化。本文只在『读』阶段(重放与答题)启用适配器,秩 $r=8$、$\alpha=16$,用不足 9.3M 通用语料 token 训练,写入与寻址完全不变。
NB-TF 与 NB-LoRA 的对比是论文机制分析的关键:1M 检索恢复完全不需要适配(59.8 vs 58.4),而 LongBench 上的大部分提升来自适配,由此区分『证据可及』与『证据利用』两类收益。
研究动机
线性注意力把全部历史压缩进一个固定尺寸循环状态,让同一状态既支撑持续计算又充当长期记忆。随着上下文增长,这两个角色互相冲突:持续压缩会在模型尚不知道未来需要什么证据之前,就不可逆地丢弃信息。已有研究证明了这一点:有限状态无法精确复制任意长序列(Jelassi et al., 2024),选择性 SSM 存在近因偏差与过平滑,多个高效注意力替代方案存在严重的联想回忆缺口。后果是灾难性的:预训练 GDN(1.3B)在 1M token 的六个 RULER-NIAH 检索任务上平均准确率仅 1.7%,Mamba-2-8B 在 MemoryAgentBench 上得 0.0 分,连『把提示读两遍』的 JRT 方法在长上下文也几乎无效。这形成了『线性注意力高效、但固定状态注定给长上下文封顶』的普遍印象。
本文的目标是本文要挑战这个看似不可避免的天花板,证明:循环状态的远程保留能力差,并不等于模型没有利用远处证据的能力;只要把证据独立保存并按需可及,冻结的预训练循环模型就能恢复强大的长上下文能力。具体目标是构造一个免训练、即插即用的稀疏情景记忆接口:在看到问题之前就把选定的文本片段归档;用主干模型自身的原生特征做写入与寻址;用主干的原生前向计算通过『重放』读取检索到的证据。整个过程不改主干权重与结构,不需要外部嵌入模型、单独训练的检索器或独立神经读取器,同时要求额外计算量与持久存储近乎可忽略。
与已有工作不同的是,已有工作的切入角度各不相同:改进状态更新规则(DeltaNet、Gated DeltaNet、Titans、TTT 等)仍在固定状态内做文章;Memory Caching、MARCH 需要在预训练阶段就缓存多个状态;RAG 依赖外部嵌入器与检索器;InfLLM 为 Transformer 检索历史 KV 块;ReadAgent 靠额外 LLM 推理选页;GritLM 需要专门的嵌入前向。NoteBank 的独特之处是『原生接口 + 文本重放』:写入用分块内原生 key 的独特性打分选段,寻址用原生 query 特征的固定坐标压缩向量,读取则把检索到的原始 token 重新喂给冻结主干走一遍普通前向。归档严格在问题出现之前完成,查询时从保留状态 $S_T$ 的副本分支检索与重放,把证据保存与循环计算彻底解耦。
核心方法
直觉上,长程检索失败的原因不是模型不会使用证据,而是证据只存在于一个被迫不断遗忘的压缩状态里。NoteBank 因此在主干的循环状态之外维护一个稀疏笔记档案:上下文按 $C$ 个 token 分块 causal 处理,每块用原生 key 的『独特性』挑出最多 $W$ 个显著句子片段存为笔记,并为每个层-头通道存一个紧凑地址向量;问题到来时,先在保留状态 $S_T$ 的临时副本上前向一遍问题,取得原生查询特征,用余弦相似度对每条笔记投票选出 top-$R$ 条;最后从 $S_T$ 的另一个副本出发,把笔记的原始 token 序列按文档顺序重放进冻结主干,再处理问题生成答案。写入、寻址、读取全程只使用冻结主干既有前向产生的特征,不引入任何可训练参数(可选的 LoRA 仅作用于读阶段)。
核心创新是『解耦 + 原生化』二者的结合。与所有改进状态更新规则的方法不同,NoteBank 根本不让循环状态承担长期存储:证据以原始 token 形式独立持久化,即使把状态重置甚至清零(Mamba-2 上 $\lambda=0$)也不丢失。与外部 RAG/InfLLM 不同,它的写与寻址复用主干在同一前向中已经算出的 key/query 特征——没有外部嵌入器、没有对比学习训练;读取则把文本重放回主干的原生前向路径——没有独立的神经读取器,任何支持并行扫描的线性注意力主干都能直接用。本质区别在于:检索不是在表示层做融合,而是让历史在查询可知之后被重新编码一次——问题到来时才决定哪些文本值得再读一遍,这正是标题『历史重现』的含义。
方法步骤详情
流程分三步(Algorithm 1)。写入:上下文按 $C{=}512$(Mamba-2 为 256)分块前向,产生各层-头 key/query;每通道计算独特性 $d_t = 1 - \langle\hat{k}_t, \hat{\mu}_{C(t)}\rangle$($\hat{\mu}$ 为块内归一化 key 均值),跨通道求和后对 48 token 滑窗打分;选最高分窗,以其中心 token 所在句为起点、在 $B{=}64$ token 预算内向两侧逐句扩展,得到最多 $W{=}2$ 条互不重叠笔记;每通道选独特性最高的位置 $t^*$,存紧凑地址 $\rho = P q_{t^*}$($P$ 为固定 $d_\rho{=}16$ 维随机坐标选择矩阵)。寻址:问题在 $S_T$ 临时副本上前向得 $u = P q_\tau$,按 $s_i = \max_\tau \cos(\hat{u}_\tau, \hat{\rho}_i)$ 让每通道为最高分笔记投票,取票数 top-$R$ 条并按原文顺序拼接。读取:从 $S_T$ 新副本重放 $S_R = F_\phi(S_{\text{read}}, \mathcal{R})$,再生成 $Y \sim p_\phi(\cdot \mid S_R, Q)$,位置索引从上下文边界延续;可选 LoRA 仅在重放与答题阶段启用。
技术新颖性
技术新颖性有四点。其一,问题无关归档:笔记在问题未知时写定,严格遵守因果设定,与 JRT 的『读两遍』提示重复有本质区别。其二,原生特征双重复用:同一次前向的 key 用于挑片段、query 用于建地址与查询,地址只是固定随机坐标投影($d_\rho{=}16$),零新增参数、零额外主干前向,寻址延迟仅 26 ms。其三,重放式读出:证据以原始 token 走完整主干(混合模型中含注意力块,按其配置的注意力策略执行),主干支持的并行算法直接可用;这与 InfLLM 复用缓存 KV 表示并在注意力内拼接的做法不同,也与 Memory Caching/MARCH 存压缩状态、需专门预训练的做法不同。其四,解耦带来新的操作自由:对状态爆炸的 Mamba-2 可以在重放前清零循环状态而证据无损,这种『重置再重放』只有在存储移出状态之后才可能实现。
实验结果
检索恢复(图3):NB-TF top-2 使 GDN 在 1M token 六个 RULER 任务均值从 1.7% 升至 70.3%,GDN/GDN-2/Mamba-2 分别超出 Plain 68.6/66.6/37.8 个百分点,且 4K–1M 无明显衰减,Plain 与 JRT 趋近于零。真实任务(表1):GDN 的 LongBench-v1 总分 15.49→16.07(NB-TF),已持平 Transformer 参考 16.05,NB-LoRA 达 23.03;Mamba-2 6.79→13.43→17.10。大模型(表2):Mamba-2-8B 在 MemoryAgentBench 0.0→14.7;Qwen3.5-9B 为 55.7→56.8、LongBench-v2 46.1→49.3。成本(图4、表7):1M 时解析 FLOPs 仅增 0.0171%–0.0250%,持久存储 74.5 MiB 约为同规格全注意力 KV 缓存的 1/2600,实测延迟增 10.6%–13.2%(125.11→140.64 s),GPU 峰值 2.93→3.27 GiB;同 checkpoint 全注意力 0 分且需 55.18 GiB,InfLLM 0 分另需 47.78 GiB CPU KV,SW4K 仅 1.17 分。诊断(图5):GDN 检索贴近归档覆盖,oracle 仅 75.0→75.7;Mamba-2 路由是瓶颈,oracle 达 40.4→61.6。路由(表3):BM25 64.0/52.7 高于原生 59.8/37.6。消融(图9/11):R 1→4 使 GDN 59.5%→75.2%,最优 $(B,R){=}(128,4)$ 时 GDN 76.8%、Mamba-2 41.3%。对照(表4):检索恢复不需 LoRA(59.8 vs 58.4),LongBench 增益主要来自适配(22.72 vs 23.03)。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| RULER-NIAH 六任务均值 @1M token(GDN 1.3B,top-2 重放) | 平均准确率 (%) | 70.3(NB-TF) | 1.7(Plain 普通前向) | +68.6 个百分点 |
| RULER-NIAH 六任务均值 @1M(GDN-2 1.3B,top-2) | 平均准确率 (%) | 约 70 | Plain 前向 | +66.6 个百分点 |
| RULER-NIAH 六任务均值 @1M(Mamba-2 1.3B,top-2) | 平均准确率 (%) | 约 38 | Plain 前向 | +37.8 个百分点 |
| LongBench-v1 12 任务(GDN 1.3B) | 宏平均总分(0-100) | 23.03(NB-LoRA)/ 16.07(NB-TF) | 15.49(Plain);16.05(Transformer 参考) | +7.54(LoRA,超过 Transformer 参考) |
| LongBench-v1 12 任务(Mamba-2 1.3B) | 宏平均总分(0-100) | 17.10(NB-LoRA)/ 13.43(NB-TF) | 6.79(Plain) | +10.31(LoRA) |
| MemoryAgentBench 2,100 题(Mamba-2-8B,top-4) | 准确率 (%) | 14.7 | 0.0(Plain);13.0(仅重置状态 Clean start) | +14.7(其中重置贡献 13.0) |
| LongBench-v2 375 题 QA/对话子集(Qwen3.5-9B,top-8) | 准确率 (%) | 49.3 | 46.1(Plain,均启用 YaRN) | +3.2 |
| 解析 FLOPs 开销 @1M token(三主干,top-2) | 相对 Plain 的增量 (%) | 0.0171–0.0250 | Plain(无增量) | 近乎零成本 |
| 持久上下文存储 @1M(GDN + NoteBank) | 字节(含循环状态) | 74.5 MiB | 约 192 GiB(同规格 BF16 全注意力 KV 缓存) | 约 1/2600 |
局限与改进
作者承认并可见于实验的局限:恢复程度依任务而异,受限重放预算下 multiquery/multivalue 等需召回多条证据的任务明显难于单针检索;Mamba-2 存在大量未兑现的路由空间(oracle 在 1M 从 40.4 提到 61.6),原生寻址在 RULER@1M 逊于 BM25(37.6 vs 52.7);Mamba-2 需在重放前清零循环状态($\lambda=0$)才能可靠读出,状态爆炸问题依旧,且清零本身贡献了大部分恢复(Clean start 已得 13.0);重放预算并非越大越好,MemoryAgentBench 上 $R$ 从 1 到 4 使 EventQA 17.9→16.4;NB-LoRA 在摘要任务退化(GDN 13.99→11.37),全文重写类任务不适合稀疏证据重放。我自己的观察:解析 FLOPs 不等于实测延迟(实测 +10.6%–13.2%);寻址是 $O(|Q| \cdot N_T \cdot J \cdot d_\rho)$ 的暴力匹配,多请求并发或超大档案需专门索引;写入显著性仅由 key 独特性决定、与未来问题无关,固定 $C/W/B$ 对任务类型无感知;评测以检索与 QA 为主,缺少仓库级代码理解等强推理任务;长上下文对照 Transformer 在该 2K 预训练 checkpoint 上得 0 分,跨架构对比需谨慎解读。
独立分析的弱点
独立分析的弱点:其一,寻址质量受限——$d_\rho{=}16$ 的随机坐标投影是无学习的固定压缩,实验显示 BM25(87 ms)在 RULER 上稳定优于原生路由(26 ms),Mamba-2 上差距达 15 分,说明原生 query 特征并非最优检索信号;改进方向是轻量学习式投影,或词法与原声特征的混合路由。其二,写入仅由 key 独特性驱动,与『未来问题需要什么』无关,可能系统性漏掉不显眼但关键的事实(例如藏在高频模板句中的数字、日期);可引入预测性保留打分或覆盖密度约束。其三,固定重放预算 $R$ 在 multiquery/multivalue 任务上是硬上限,可按查询复杂度自适应调整 $R$ 或做迭代式检索-重放。其四,状态清零是对 Mamba-2 症状的缓解而非治疗,根因是状态幅度随长度爆炸,值得研究状态归一化或正则化。其五,LongBench 增益大多来自 LoRA 而非记忆本身(有 NB 23.03 vs 无 NB 22.72),说明冻结模型利用重放证据的能力有限,读阶段仍有很大提升空间。
未来方向
作者提出的方向:对写入、路由、读取做选择性训练,分别提升证据保存、访问与利用;改进路由是明确的头寸(oracle 与 BM25 都显示空间)。基于成果可延伸的方向包括:把归档-重放与预训练结合,例如在 16K+ 长上下文训练的混合模型上协同优化;面向智能体的跨会话流式记忆系统——档案天然适合持久化多轮交互历史(MemoryAgentBench 方向);把 BM25 等词法索引与原生投票融合并控制在低延迟内;对 Mamba-2 状态动力学的理论分析与稳定化,使重放不必依赖清零;推广到 xLSTM、TTT、Titans 等其他循环主干;自适应重放预算与多跳问题的顺序多次重放;以及把『归档原始文本』扩展为『归档中间推理结果或结构化记忆』,让长程多步推理任务也能受益于重放。
复现评估
复现条件良好。附录给出完整超参与会计:$C{=}512/256$、$W{=}2$、$B{=}64$、$d_\rho{=}16$、48 token 打分窗、32 token 滚动尾巴、$R\in\{1,2,4\}$;地址构建与 FLOPs/存储逐项公式;固定评测样本(RULER 每任务 100 例、LongBench-v1 每子集 30 例、MemoryAgentBench 共 2,100 题);LoRA 细节($r{=}8$、$\alpha{=}16$、600 步 AdamW、学习率 $5\times10^{-4}$、不足 9.3M token 通用语料)。所有主干均为公开 checkpoint(m-a-p/1.3B-100B-GatedDeltaNet-pure、state-spaces/mamba2-1.3b、nvidia/mamba2-8b-3t-4k、Qwen/Qwen3.5-9B 等),1.3B 模型单张 H100 即可做 1M token 推理(峰值含权重约 3.3 GiB)。项目页 nanboli.github.io/notebank-page 提供概览,但正文未明确声明开源完整代码仓库,实现需基于 FLA/mamba 生态自行组装,Mamba-2 还需实现状态重置技巧。总体难度中等:免训练部分概念简单、算力门槛低,长上下文评测基础设施与 LoRA 训练流水线需要一定工程投入。
论文图表
左图展示三个 1.3B 预训练循环模型(GDN、GDN-2、Mamba-2)在六个 RULER-NIAH 检索任务上从 4K 到 1M token 的平均准确率,虚线为无 NoteBank 的 Plain,实线为免训练 top-2 重放;右图对比 1M token 时的解析 FLOPs 增量(≤0.025%)与持久存储:GDN+NoteBank 共 74.5 MiB,约为同规格全注意力 192 GiB KV 缓存的 1/2600。
这是全文的概览图:一张图同时给出『恢复效果』与『近乎零成本』两大核心主张,是理解论文贡献定位的入口。
GDN 在有无 NoteBank × 冻结/LoRA 读取器的 2×2 对照:RULER@1M 上无 NB 时 LoRA 仍只有 1.7,有 NB 冻结即 59.8(LoRA 58.4);LongBench 上无 NB LoRA 22.72、有 NB LoRA 23.03,仅差 0.31。
这是机制层面最关键的对照:证明百万 token 检索恢复完全来自记忆访问而非参数适配,而 LongBench 增益主要来自适配,避免混淆两种效应。