Paper

论文速读:让模型自己决定何时回看——On-Demand Attention 用一个小 recall head 替掉全量注意力

6 min read ·

长上下文推理的账本上有一笔长期没被认真审的支出:你为「可能用得上」付费,而不是为「真的用得上」付费。全量注意力解码在每一步都把整段历史的键值读一遍,不管这一步的预测究竟依赖多远的信息。序列越长,这笔支出越大,而且它是固定的——不随内容变化。

On-Demand Attention(下称 ODA)这篇论文做的事情,就是把这笔固定支出变成浮动支出。

一、先回答一个前置问题:模型自己知道该不该回看吗

任何「跳过计算」的方案都绕不开一个前置问题:你靠什么判断这次可以跳过。规则、启发式、外部检索器都可以,但它们要么不通用,要么本身就是额外开销。

论文的切入点是观察而不是假设。作者发现,预训练模型在执行全局读取之前的解码状态里,已经含有能够预测「这次读取会带来多少收益」的信息。这句话的分量在工程上很实在:

  • 判断不需要额外的前向计算,只要把已有状态读出来;
  • 判断不需要外部规则,所以能跟着模型能力一起演进;
  • 判断发生在读取之前,因此有资格决定要不要读。

很多稀疏化方案的问题在于「先读再决定要不要留」,顺序反了,省不下带宽。ODA 的顺序是对的。

二、方法:局部优先解码加一个按需触发的全局分支

把上面的观察变成结构,ODA 的设计可以拆成三层。

第一层是默认路径。 解码默认走局部优先(local-first):只看最近的一段窗口。这是成本最低、也符合「多数 token 只依赖近处」这一经验事实的路径。

第二层是收益预测。 一个轻量的 recall head 读取当前解码状态,输出「现在做一次全局读取值不值得」的预测。这个预测在生成过程中是连续变化的——同一个序列,有的位置收益高,有的位置接近零。

第三层是按需触发。 当预测收益越过阈值时,才真正发起全局注意力;否则继续走局部路径。

这里有两个设计选择值得单独说,因为它们决定了方案能不能落地。

其一,只训练 recall head,预训练权重全部冻结。 冻结主干带来三个直接好处:原模型的通用能力不被破坏,不存在定向微调造成的退化风险;不需要重新做对齐与安全评估;训练成本比全参数微调低好几个数量级,团队可以在自己的模型和数据上快速适配。这是一种很典型的「冻结主干、加一个决策适配器」范式——学的是策略,不是知识。

其二,完整历史 KV cache 依然保留。 这一点常被误解。局部优先不等于丢弃:历史一直躺在显存里,随时可被召回。这和第 4 步才决定读不读的逻辑是自洽的——如果提前把远处历史裁掉,后面的「按需」就成了空话。

三、从方法到可部署:为什么必须动 vLLM 的执行路径

论文在 vLLM 里实现了 GPU 侧的条件执行,这是我认为最能体现工程判断的一处。

想象另一种做法:在算法层判断出「这次不需要全局读取」,然后把分支带回 CPU 侧决策,再决定是否发射内核。听起来省了读取,实际会怎样?同步开销、调度延迟、内核发射成本会层层累加,最后端到端测不出任何加速,甚至更慢。长上下文场景下这类「论文里省了 90%,实机测下来省了 3%」的例子并不少见。

把条件执行放在 GPU 侧,让分支判断和设备上的执行保持在同一层,减少的全局读取才能真正转化为解码加速。这也说明一件事:推理优化的收益往往不是被算法本身卡住的,而是被执行路径的层级错配卡住的。

四、实验覆盖了什么

论文在 Qwen 与 Gemma 系列模型上做了验证,并且特意包含了混合注意力主干。这个选择很关键,因为混合架构(在层间交错注意力和循环层)正在成为长上下文模型的主流形态,而循环层的结构与注意力不同,按需召回在这里是否依然成立,不能从纯注意力模型的结果外推。

论文报告的结论是两条:

  1. 选择性召回恢复了局部注意力下损失的大部分性能——也就是说,按需回看能拿回「全量读」的绝大多数好处;
  2. 同时显著减少全局读取次数——省的是真实的带宽与算力,而不是账面上的理论值。

全文 28 页、5 张图,提交于 2026 年 9 月 17 日。需要注意的是,摘要层面的表述是「大部分性能恢复」而不是「无损」,这在稀疏化这类方案里是常态,选型时必须按自己的任务分布实测,而不是照搬结论。

五、放在更大的图景里看

如果把长上下文优化拆成几个正交层面,位置会很清晰:

层面改什么代表工作效果
架构层哪些层读全局混合注意力、层间交错循环层决定可读范围
存储层历史以什么精度存KV cache 量化与熵编码决定占用与带宽上限
访问层每一步读不读On-Demand Attention决定实际读取时机
调度层怎么并行执行分页注意力、条件执行内核决定收益能否兑现

这几层并不互斥。架构决定「最多能看到多远」,存储决定「看到要花多少内存」,访问策略决定「这次要不要看」,调度决定「省下来的操作会不会被开销吃掉」。ODA 补的是第三层,而且它选择了一个很聪明的训练预算——只学策略,不学知识。

同一批 9 月中旬的论文里,还有一篇 dQwen3.5 走的是另一条路:把预训练的自回归混合主干(交错注意力与循环层)适配成扩散语言模型,报告在达到同等训练损失时只用了约一半 token。两者的共同点值得留意——都在从「把主干做大」转向「把已有主干用得更省」。这个方向的变化,对硬件预算有限的团队是明确的利好。

六、如果你要把它用起来

  1. 先测量收益分布:在自己业务的真实轨迹上统计「需要远距离信息」的 token 占比。如果这个比例很低,按需召回的空间就大;如果任务本身就是长程依赖密集型的,收益会明显缩水。
  2. 把阈值当超参而不是常量:预测收益的触发阈值直接决定「省了多少」与「丢了多少」的权衡点,应按任务调。
  3. 优先验证执行路径:落地前先确认你的服务栈支持条件执行。算法层省下的读取,如果在调度层被吃掉,等于没做。
  4. 用适配器隔离风险:只训一个头意味着随时可以关掉它回到全量注意力,这是很宝贵的回滚能力——先灰度,再放量。
  5. 和自己的压缩方案组合:访问策略与 KV 量化并不冲突,两者叠加往往比单独用任一个更好。

这篇论文真正值得记住的,可能不是那两个具体模块,而是它提出的那个顺序问题:如果你能在读取之前就知道这次读取值不值,那么「读多少」就从一个架构常量变成了一个运行时可调的策略。

Frequently asked questions

为什么说全量注意力解码存在浪费?
因为它在每一步都假设「所有历史都可能有用」。但实际生成过程里,大部分 token 的预测只依赖最近的一小段上下文,只有少数位置才需要回头去取很久以前的信息。全量注意力不管这个差别,每生成一个 token 就把整段历史的键值读一遍,代价随序列长度线性增长,还要加上带宽开销。浪费不在于读得多,而在于读得不加区分——把确定有用的近处信息和大概率无关的远处信息按同样权重处理。
「解码状态已经含有收益信号」具体指什么?
指在真正执行全局读取之前,模型当前的隐状态里就已经编码了「这次回看是否必要」的信息。这是一个很有分量的观察,因为它意味着一件事:判断要不要回看,不需要额外跑一次前向计算,也不需要外部规则,直接从模型已有的状态里读出来即可。这也解释了为什么论文能把新增训练量压到极小——需要学的不是「怎么理解长上下文」,而只是「怎么把已有状态里的这个信号读出来」。
只训练 recall head 有什么工程价值?
三点。第一,预训练权重完全冻结,所以原模型的通用能力不会被破坏,也不用担心定向微调带来的退化。第二,完整历史 KV cache 依然保留,意味着后续任何时候都可以把这个 token 召回回来,局部优先不等于永久丢弃。第三,训练成本极低,一个轻量的头相比全参数微调差了好几个数量级,团队可以在自己的模型和数据上快速适配。这种「冻结主干、只加一个决策适配器」的范式在推理优化里会越来越常见。
为什么要在 vLLM 里做 GPU 侧条件执行?
因为只做算法优化不够。如果判断出「这次不需要全局读取」之后,还要把控制流里的分支带回 CPU 侧决策、再决定是否发射内核,节省的读取量就会被同步与调度开销吃掉,最终测不出端到端的收益。把条件执行放到 GPU 侧,让分支判断和设备上的执行保持在同一层,减少的全局读取才能真正转化为解码加速。这一步是论文从方法变成可部署方案的关键。
它和混合注意力、KV 压缩是替代关系吗?
更像是正交关系,可以叠加。混合注意力改的是「模型里哪些层是局部、哪些是全局」这件事,属于架构层的静态设计;KV 压缩改的是「历史在内存里以什么精度存放」,属于存储层的表示问题;On-Demand Attention 改的是「每一步到底读不读」,属于访问策略层。三者作用于不同层面:架构决定可读范围,压缩决定读取成本,访问策略决定读取时机。叠加起来才能同时改善容量、带宽和算力三个约束。
// next.txt ›

Some outbound links in this post are affiliate links — see disclosure.