Back to Blog 模型自己管注意力:一句话声明,KV cache 读取量减半

模型自己管注意力:一句话声明,KV cache 读取量减半

Paper
Language Models Can Control Their Own Attention
KAIST AI · Google DeepMind · arXiv 2609.02737 · 2026-09-02
稀疏注意力KV Cache长上下文零训练
一句话:让模型在思维链里用 <global> / <focus> / <local> 三种标签声明自己要 attend 哪里,推理引擎像解析工具调用一样解析声明、构造段级注意力掩码、跳过大部分 KV cache 读取。零训练、零辅助打分器,现成模型上注意 token 减少 52.0% / 31.1%,准确率只掉 1.27pp / 2.75pp。

一、一个被忽视的常识

长上下文推理最贵的不是算力,是搬数据。解码每生成一个 token,都要把整个 KV cache 从显存读一遍——Qwen-3.5-397B 在 1M 上下文下,每个序列每步要读约 15 GB KV cache,这个带宽需求跟加载 17B 激活参数一个量级。而实证研究早就发现:模型真正"在意"的只有上下文里一小撮 token,却每次都全量扫描。

现有稀疏注意力方案的共同套路是外挂打分器:用轻量代理分数预选相关 token,再做 attention。问题是打分本身每步仍是 O(N)——你只是把全量注意力换成了全量打分,账没省掉多少。

这篇论文的出发点是个朴素的反问:模型自己难道不知道该看哪吗?

二、Declarative Attention:把注意力变成声明式

答案是:知道,而且能直接说出来。作者设计了一套协议,让模型在思维链里声明自己的注意力去向,三种模式:

<global> 浏览全上下文,用于导航定位;<focus chunk=N> 只看编号 N 的"魔法块"(长上下文预先切成的段落);<local> 只看自己已生成的回答,用于收敛。

method
图 1 完整机制:25,466 token 的 prompt 被切成 12 个魔法块,模型先 <global> 定位(全量注意),再 <focus chunk=1> 精读(−86.5%),最后 <local> 收敛(−95.6%)。引擎从文本轨迹直接派生段级掩码,无需任何辅助打分器。

推理引擎像解析工具调用一样解析这些标签,动态构造 segment 级注意力掩码——<focus> 时其余块直接从 KV 读取中跳过,没有任何打分步骤。不需要训练,不需要辅助模型,现成模型零样本可用。

协议本身的设计相当精细:每次回答至少一个 <focus>(最常见的失败模式是跳过定位直接靠记忆猜);必须以 <local> 收尾并点名最终答案("承诺步骤",防止取回一堆值后忘了题目要哪个);进入 <local> 后魔法块已被掩蔽,"记忆里"的块内容都算瞎猜,需要就回去再 <global>

三、效果:零样本就是下限

指标Gemma-4-31BQwen-3.6-27B
注意 token 减少52.0%31.1%
准确率损失(15 任务均值)−1.27 pp−2.75 pp
估算解码时延(B200 roofline)0.71×0.77×
KV 读取量(整个生成过程)1100 → 528 GB1477 → 1017 GB
acc
图 2a 相对准确率:DA 掩码承担了主要损失,DA-nm(仅块状格式、无掩码)几乎无损
tok
图 2c 相对注意 token:掩码贡献了几乎全部节省

三个细节值得咀嚼:

省的 token 比省的时延多。 DA 会多生成 31–35% 的推理文本(声明本身有开销),matmul 和本地内存成本反而微涨;Gemma 的 SWA 层(60 层里占 50 层)是不可压缩的读取下限,占住 42% 的注意力时间,把节省封了顶。Qwen 的 GDN 循环状态很小(5%),节省几乎全额传导——这就是 0.71× 和 0.77× 差异的来源。注意这是 roofline 理论估算,非实测。

越强的模型协议遵守越好。 focus 成功率从 Gemma-4-E4B 的 58% 升到 31B 的 99%;相对准确率沿家族规模单调爬升(29%→99%,64%→97%)。DA 本质上是一项"能力活":既要解析协议、又要克制自己不全量阅读,两者都吃模型智力。

scaling
图 3 规模缩放:两个家族内相对准确率随模型规模单调提升,Gemma-4-31B 达到 vanilla 的 99%

模式分布也印证了协议意图:<global> 只占生成 token 的约 27%,73% 的推理发生在廉价的 <focus>/<local> 里,后两者每 token 注意力节省 76–99%。

modeshare
图 5a 模式占比随上下文长度变化:最长上下文桶里 <global> 回升到约 45%
saved
图 5b 每 token 注意力节省:<focus>/<local> 节省 76–99%,且随上下文变长而扩大
adherence
图 6 协议遵守率随模型规模单调上升:Gemma-4-E4B 58% → 31B 99%,Qwen-3.5-4B 89% → 27B 99%

四、为什么值得认真对待

第一,它开了一条正交的稀疏化轴。 KV 量化压的是"每个 token 读几个字节",DA 压的是"读多少 token"——两者是乘法关系,可以堆叠。KV 读取量已经减半的模型再叠上量化,长上下文解码的经济性会完全改写。

第二,"声明式"这个范式比数字本身重要。 过去的稀疏注意力是让引擎去模型要什么(隐藏状态探针、代理分数),DA 是让模型它要什么。这是 tool use 范式向推理内核的渗透:模型不只调用外部工具,还"调用"自己的注意力。声明写在文本轨迹里,可审计、可解释、可干预——对安全审查也是利好。

第三,对视频语义检索是同构启发。 少量关键帧含目标信息时,全局 softmax 的竞争会稀释关键信号——这正是筛选机制天然适配该场景的原因。DA 相当于把筛选权交给模型自己:先 <global> 扫块级摘要,再 <focus> 精读候选段,粒度是段而非帧。零训练就能跑,意味着这个范式可以直接嫁接到多模态底座上试。

五、冷水

消融显示掩码贡献了几乎全部 token 节省,也贡献了大部分准确率损失(Qwen 上 −2.75pp 里 −2.06pp 来自掩码,块状格式本身几乎免费);长上下文桶里 <global> 占比回升到 45%,节省被侵蚀;少数响应(Gemma-4-12B 上约 6%)在 8K 预算内收不了尾。零样本结果是下限——论文自己也承认,训练加持下的 DA 才是完全体。

一句话总结:与其让引擎猜模型的心思,不如让模型亲口说出来。这可能是稀疏注意力里最"agent 味"的一条路线。
Tags: #Blog