← 返回 2026-09-03

语言模型可以控制自身的注意力 Language Models Can Control Their Own Attention

Namgyu Ho, Huzama Ahmad, Woosung Koh, Se-Young Yun, Tal Schuster, Cicero Nogueira dos Santos 📅 2026-09-02 👍 77 2026-11-02 18:30
KV缓存 推理优化 提示工程 稀疏注意力 长上下文

让模型在思维链中声明注意力范围,推理引擎据此跳过大部分KV缓存读取

前置知识

KV缓存(Key-Value Cache)

Transformer 自回归生成时,每个已生成的 token 在每层都会产生一对 Key/Value 向量,为避免重复计算会把它们全部缓存下来。解码的每一步都要把这个缓存从显存(HBM)完整读一遍供注意力查询。上下文越长缓存越大,例如论文测得 Qwen-3.5-397B-A17B 在 1M token 上下文下每个序列每步要读约 15GB 的 KV 缓存,与加载 17B 激活参数的带宽需求相当。

本文所有收益都来自跳过 KV 缓存的读取,理解了它是解码延迟的主导项,才能理解为什么『少读缓存』比『少算 FLOPs』更值钱。

自回归解码与注意力读取

解码是逐 token 进行的:每生成一个 token,所有全局注意力层都要对全部历史位置做一次注意力计算。即便注意力权重实际集中在少数 token 上,硬件层面仍要读完整 KV 缓存才知道该关注谁,因为真实注意力分数只有在算完整个注意力矩阵后才可知。

论文的核心矛盾就藏在这里:想先知道该读哪里,就必须先读完;这解释了为什么以往方法只能用代理分近似,而 DA 换了思路。

稀疏注意力(Sparse Attention)

只让每个查询位置关注一部分 key/value 而非全部上下文的机制。静态方案(滑动窗口、注意力汇聚点)按固定规则裁剪;动态方案如 Quest、DeepSeek 的轻量索引器在每步用代理分数扫描 KV 缓存挑出高分 token。前者无法预知未来查询需要什么,后者每步仍是 $O(N)$ 的扫描开销。

DA 是与这两类都正交的第三条路:让模型用文字声明注意力范围,把每步 $O(N)$ 的选择成本彻底归零,论文反复与这两类基线对比。

Roofline 性能模型

判断一个算子受算力还是受内存带宽限制的经典模型。计算受限时耗时 $T_{\text{FFN}} = \text{FLOPs}/(\text{峰值FLOPS}\times\text{MFU})$,带宽受限时 $T_{\text{attn}} = \text{KV字节数}/(\text{峰值带宽}\times\text{MBU})$,其中 MFU/MBU 是实际达到的峰值利用率。大 batch 解码时 FFN 已摊薄成计算受限,而注意力 KV 读取仍是带宽受限。

论文用它论证『注意力何时主导解码成本、DA 何时划算』,并把节省的 token 数换算成 0.71×/0.77× 的理论墙钟时间,读懂它能看懂第 3、5.4 节。

思维链(CoT)与工具调用格式

思维链让模型把中间推理写成可读文本再给最终答案,研究表明隐状态中已编码了对未来 token 的预测,CoT 只是把这种潜在计算显式化。工具调用则是后训练对话的常见格式,user/assistant/tool 消息之间有模型非常熟悉的特殊 token 边界。

DA 把 CoT 从『决定想什么』扩展到『决定看哪里』,并把长上下文伪装成模型『已经调用过的工具返回结果』(magic chunks),正是利用模型对这两种格式的熟练度实现零样本生效。

研究动机

Transformer 解码的每一步都要对所有历史位置做注意力,导致必须从 HBM 完整读取 KV 缓存。论文给出具体量级:Qwen-3.5-397B-A17B 在 1M token 上下文下,每个序列每个解码步要加载约 15GB KV 缓存,带宽需求与加载其 17B 激活参数相当——长上下文场景下 KV 读取而非计算才是解码延迟的支配项。经验研究表明真实注意力权重本来就高度集中在少数 token 上,于是稀疏化成了自然解法,但存在根本困难:真实注意力分数只有在算完完整注意力矩阵后才可知。早期方案用静态启发式(近因窗口、历史注意力强度)预测该保留哪些 token,但固定规则无法预知未来查询需要什么,长上下文性能明显退化;近期方案(如 Quest 类轻量扫描、DeepSeek 的索引器)承认注意力依赖查询,在每步用轻量代理分扫描 KV 缓存选出高分子集,然而这只是降低了常数因子,每步复杂度仍是 $O(N)$——在 1M token 的对话里,每生成一个回复 token 都要做一次全局扫描。

本文的目标是本文要回答的问题是:模型自己难道不知道上下文里哪些部分相关吗?目标是设计一个任务无关、无需任何训练的协议(Declarative Attention, DA),诱导现成模型在思维链中逐步『声明』自己接下来要关注上下文的哪个区域,推理引擎直接从模型生成的文本解析出注意力掩码并据此跳过大部分 KV 缓存读取,从而彻底消除每步 $O(N)$ 的 token 选择开销——$O(N)$ 的完整读取只发生在模型显式声明的 global 阶段。同时系统量化这套协议在精度、注意力成本、理论墙钟时间三个维度的权衡。

与已有工作不同的是,本文的独特切入是『内在声明』而非『外在评分』。以往工作都在模型外部构造代理分数来近似注意力掩码,DA 则把掩码的来源换成了模型自己的输出文本:既然 CoT 能把模型对『想什么』的隐式计算显式化,那也可以让它把『看哪里』说出来。与最接近的先前工作 Self-Selected Attention(Jin et al., 2024)相比,后者需要按任务分别微调、在手工设计的上下文划分和标注语法下只工作于 2K token 上下文;DA 证明用单一固定提示词就能在现成模型上零样本诱导出同样行为,且上下文规模达 100K token 量级。由于掩码直接读自人类可读的推理轨迹,注意力计划天然可审计——驱动 KV 读取的 token 同时就是让人能监督它的 token,这是从隐状态近似掩码的方法做不到的。

核心方法

先说直觉:让模型像人读长文档一样先浏览、再精读、后综合,并把『我现在在看哪里』写成文本,推理引擎照着这份声明跳读缓存。技术路线分三层。第一层是协议:DA 把生成过程划分为三种模式—— 浏览全部上下文做导航, 只关注指名的特定区域做精读, 只看自己已生成的输出做推演和汇总,模型可以在 CoT 中任意交替使用。第二层是上下文投递:把长输入切成约 2048 token 的可寻址片段,称为 magic chunks,并用模拟工具调用 transcript 呈现(assistant 轮调用 get_magic_chunk、tool 轮返回片段),使段边界落在模型在后训练中烂熟的对话特殊 token 上。第三层是解码时干预:一个与推理引擎并行的 DA 状态机监视输出流、解析标签切换、构造块对齐注意力掩码,通过 vLLM 钩子逐步重写请求的 KV block table,无需改动任何注意力内核。

核心创新是把注意力掩码的来源从『隐状态近似』换成『模型生成的文本』,这是与所有已有稀疏注意力方法的本质区别。轻量扫描方法每步都要为每个上下文 token 打分($O(N)$),DA 的声明则让引擎在段级别跳读:focus 模式只读指名 chunk、local 模式完全不读上下文,两者每 token 节省 76–99% 的注意力读取。掩码以块粒度对齐(16–32 token 的 vLLM 块,保留区间向外取整),得到的是普通块列表,FlashAttention 等现有内核原样运行,沿用 Native Sparse Attention 的块稀疏原则。DA 只作用于全局注意力层,SWA/GDN 等高效层每步成本有固定上限、无可节省故不动。作者把这称为 system-2 稀疏注意力:选择策略写在语言里,可随指令改变、随模型规模自增强,而且序列从不被编辑、被遮蔽内容的缓存始终有效,为 KV 卸载和可逆上下文压缩留出接口。

方法步骤详情

第一步,分段:目标 2048 token,只切超限单元,按最粗边界优先(段落→换行→句末→子句→词边界),段边不落词内。第二步,组装提示:恒定支架含系统指令(占据 attention sink)、用户问题与三模式指引;上下文区为伪造的 get_magic_chunk 工具 transcript。第三步,生成:模型在 CoT 中交替输出 、、,末尾用 作答。第四步,状态机:默认 global,读开标签的闭合 > 切换,读闭合标签回退。第五步,块对齐掩码:保留区间向外取整到 16–32 token 的 KV 块,每步重写 vLLM 的 block table。第六步,评估:15 来源、每源≤128 例;Gemini-3-Flash 生成 rubric、thinking Qwen-3.5-4B 评审(Pearson $r=0.99$);B200 + vLLM、上限 8K。

技术新颖性

新颖性有四点。其一,选择成本归零:代理分方法每步 $O(N)$,DA 的 $O(N)$ 完整读取只剩模型自己声明的 global 阶段,不再是逐步开销。其二,零样本即可用:不训练任何参数,固定一套提示词就跨 15 个任务生效,且论文明确这是下界——若针对协议做后训练还有大量空间。其三,系统级可行:块对齐掩码设计让 DA 无需改内核就能搭上 FlashAttention 的效率快车,掩码只在模式边界变化,天然适配投机解码(草稿与验证共用单一掩码)和 KV 卸载(声明即预取信号)。其四,给出了何时划算的定量框架:用 roofline 分析证明收益上界由全局注意力在解码时间中的占比决定(本文两个主模型为 73% 和 86%),并指出在最激进的高效架构(MLA、GDN、索引器)上该占比仍达 56–97%,问题不会很快消失。

Declarative Attention (DA) lets a model control its own attention span.
Figure 1: Declarative Attention (DA) lets a model control its own attention span.

实验结果

15 个任务平均:Gemma-4-31B 准确率 87.01%→85.74%(−1.27pp),attended tokens 13.43M→6.45M(−52.0%);Qwen-3.6-27B 准确率 85.31%→82.56%(−2.75pp),tokens 22.54M→15.52M(−31.1%)。Gemma 在 7/15 任务持平或超 vanilla,最大增益为 longdep_qa(+3.1pp)与 code_repo(+5.6pp)。消融显示分块格式几乎无损但 DAnm 反而多注意 66.2%/28.8% 的 token,掩码相对 DAnm 再砍 71.1%/46.5%:效率与精度损失均主要来自掩码。相对准确率随规模单调上升(E4B 29%→31B 99%;Qwen 4B 64%→27B 97%)。32K 内精度差 ≤1pp、最长桶约 96%,绝对节省从 −1M 增至 −21M token/响应。focus+local 每 token 省 76–99%,global 占生成 27% 却占超 80% 的 attended tokens,墙钟估计降至 0.71×/0.77×。

The benchmark suite spans 15 long-context sources, from short documents to million-token code repositories.
Table 1: The benchmark suite spans 15 long-context sources, from short documents to million-token code repositories.
DA reduces attention cost with a 1–3pp average accuracy drop.
Table 2: DA reduces attention cost with a 1–3pp average accuracy drop.
Estimated decode wall-clock time on an optimized disaggregated inference setup.
Table 3: Estimated decode wall-clock time on an optimized disaggregated inference setup.
DA's efficiency comes from attention masking rather than shorter generation.
Figure 2: DA's efficiency comes from attention masking rather than shorter generation.
DA accuracy scales favorably with backbone size.
Figure 3: DA accuracy scales favorably with backbone size.
DA holds accuracy near vanilla while its absolute token saving grows with context length.
Figure 4: DA holds accuracy near vanilla while its absolute token saving grows with context length.
DA spends most tokens in cheap modes, with larger savings at longer contexts.
Figure 5: DA spends most tokens in cheap modes, with larger savings at longer contexts.
DA focus adherence improves with model size.
Figure 6: DA focus adherence improves with model size.
查看结构化数据
任务指标本文基线提升
15 个长上下文任务总体(Gemma-4-31B) 平均准确率(LLM judge) DA 85.74% Vanilla 87.01% −1.27pp,同时 attended tokens 减少 52.0%(13.43M→6.45M)
15 个长上下文任务总体(Qwen-3.6-27B) 平均准确率(LLM judge) DA 82.56% Vanilla 85.31% −2.75pp,同时 attended tokens 减少 31.1%(22.54M→15.52M)
LBv2/code_repo(百万 token 级代码库问答,Qwen-3.6-27B) 准确率 DA 77.8% Vanilla 72.2% +5.6pp,attended tokens 134.00M→81.96M(−38.8%)
LooGLE/longdep_qa(长依赖问答,Gemma-4-31B) 准确率 DA 68.8% Vanilla 65.6% +3.1pp,attended tokens 10.57M→6.54M(−38.1%)
解码墙钟时间(roofline 估计,B200,MFU 40%/MBU 70%) 相对 vanilla 的每响应解码耗时 Gemma-4-31B 0.71×(269.1ms→192.3ms) Vanilla 1.00×(Gemma 269.1ms;Qwen 306.2ms) Qwen-3.6-27B 为 0.77×(306.2ms→237.3ms)
协议遵循(<focus> 标签解析为有效 chunk 引用的比例) Focus 成功率 Gemma-4-31B 99%、Qwen-3.6-27B 99% Gemma-4-E4B 仅 58% 随模型规模单调上升,说明小模型瓶颈在遵循协议而非推理

局限与改进

作者承认的局限:全部结果是零样本提示的产物而非机制上限——DA 比 vanilla 多跑 15–35% 解码步、模式使用策略未必贴合任务;thinking 模式下协议失效只能关闭评估;固定基准迫使人为切块;global 模式占 attended tokens 超过 80% 且随上下文上升;墙钟时间是 roofline 理论估计而非实测。我的补充:多跨推理任务掉分明显大于单跨(Gemma 2.28pp vs 0.78pp,Qwen 3.59pp vs 2.34pp),信息分散时声明式跳读更伤;精度损失与掩码强绑定,本质是可控的精度换读取;小模型协议崩溃(E4B focus 成功率仅 58%)构成最低能力门槛;评估依赖 LLM judge、每源最多 128 样本、仅覆盖 Gemma/Qwen 两家族;另有不到 1%(Gemma-4-12B 约 6%)的响应在 8K 预算内不终止,会污染 attended tokens 统计。

独立分析的弱点

第一,global 模式是成本黑洞:仅占 27% 生成 token 却贡献超 80% 的 attended tokens,份额随上下文升至 45%,长上下文下总节省被封顶。改进方向是作者提出的上下文索引(每 chunk 配极短描述,导航只读索引),或在 global 步叠加轻量扫描稀疏注意力。第二,零样本策略非最优:DA 多跑 15–35% 解码步,模式切换由模型即兴决定,与经 SFT+RL 深度优化的 vanilla CoT 不对等,应对 DA 轨迹做后训练、用同时奖励精度与效率的 RL 优化。第三,thinking 模式失效把方法锁在非思考场景,而最长推理轨迹恰在思考链中;把 DA 三种操作暴露为标准工具声明是自然出路。第四,roofline 墙钟只是理论值,缺少端到端实测;掩码同时是精度损失来源,多跨任务掉分更多,需要 focus 验证或失败自动回退 global 的兜底机制。第五,协议遵循的规模门槛把 4B 以下模型排除在外,可通过蒸馏 DA 轨迹或小规模协议微调降低门槛。

未来方向

作者提出的方向:一是对 DA 做后训练(SFT+RL),像优化推理链一样优化注意力声明策略,奖励精度与效率;二是与轻量扫描稀疏注意力协同——用扫描方法服务 global 步、声明负责其余模式,且 DA 通过改写块表施加掩码,集成代价低;三是与投机解码叠加——DA 降低每步成本、投机减少步数,验证趟也只读掩码后的 KV,两项收益可乘;四是 system-2 KV cache 卸载:模式边界处的声明是可预取的明确信号,失焦段可溢出到主存并在被点名前搬回,进而实现可逆上下文压缩——序列不编辑、缓存不失效,无需重新 prefill 即可恢复被压缩内容;五是在 agentic 场景用天然段(工具返回、对话轮次)替代人工切块。基于成果还可延伸:把上下文索引进 global 导航、用 DA 轨迹作为可解释性审计界面研究注意力行为、评估 DA 在 1M+ 真实长程 agent 任务上的表现,以及与 MLA/低秩 KV 等压缩表示结合后重新核算收益边界。

复现评估

复现材料较为齐备:模型均为公开权重(Gemma-4 系列、Qwen-3.5/3.6 系列),基准全部公开(RULER、LongBench v1/v2、LooGLE、ZeroSCROLLS),论文附录给出了完整的 DA 指令与 vanilla 对照提示词(附录 F)、分段器细节、vLLM 集成方式(附录 B)和 roofline 推导(附录 C),方法本身无需训练。主要门槛在算力与工程:主实验需要 NVIDIA B200 级 GPU 集群承载 27B–31B 模型的 244K 上下文推理,自定义 vLLM 的注意力元数据钩子需要一定系统工程能力;部分合成 QA 依赖 Gemini-3-Flash 生成,评审依赖 thinking 模式的 Qwen-3.5-4B 并以 Gemini-3.1-Pro 校验(Pearson $r=0.99$),会引入 API 成本与评审方差;正文未见明确的开源代码声明,若不公开则 block table 重写等细节需按附录自行实现。总体评估:有 B200 资源可复现主表量级,roofline 分析部分只需模型架构规格即可独立核算,整体复现难度中等偏高。

相似论文推荐