语言模型可以控制自己的注意力
语言模型将大部分注意力集中于少数关键上下文,却需要读取整个 KV 缓存来定位这些 token。在百万级 token 的对话中,全局注意力层为生成每个回复 token,都不得不扫描全部上下文。现有缓解方法通过轻量级代理分数预选相关 token,但这种外部打分仍带来每步 O(N) 开销。 我们提出 声明式注意力(Declarative Attention, DA)协议,让模型在思维链中主动声明需要聚焦的上下文区域,并将生成过程划分为 global(全文)、focus(特定区域)和 local(仅最近输出)三种模式。推理引擎像解析工具调用一样处理这些声明,从而跳过大部分 KV 缓存读取。 在 15 个长上下文任务的零样本评测中,DA 应用于现成模型(Gemma-4-31B、Qwen-3.6-27B),显著减少解码期间总关注 token 数(52.0%、31.1%),同时仅带来 1.27pp 与 2.75pp 的精度下降,且随模型规模增大差距进一步缩小。DA 为稀疏注意力开辟了新方向,未来结合训练类方法有望进一步释放潜力。
论文精读
TL;DR Declarative Attention 让模型在链式思考中主动声明注意力范围(<global>/<focus>/<local>),推理引擎据此跳过大部分 KV cache 读取,在长上下文解码中大幅减少注意力 token 数,精度损失随模型规模增大而缩小。
问题
问题背景
长上下文语言模型在解码时,全局注意力层必须扫描完整 KV cache 才能生成每个 token。即使模型最终只关注少量相关 token,计算和内存带宽开销仍随上下文长度线性增长。
现有方法局限
主流方案通过轻量代理分数(proxy scores)预选相关 token,例如基于注意力统计或检索的稀疏注意力。但这些方法存在明显不足:
- 每步仍需计算所有 token 的分数,复杂度为 O(N),并未真正消除全局扫描;
- 代理分数与主模型分布存在偏差,可能误选或漏选关键 token,影响生成质量;
- 额外存储和读取代理分数消耗内存带宽,在 1M-token 级上下文中收益有限。
为什么这个问题难且重要
核心挑战在于模型自身在推理过程中其实“知道”哪些上下文片段与当前生成相关,但缺乏一种低成本、内生的表达机制。如果能让模型直接声明关注区域,就能从根本上跳过无关 KV cache 读取,将解码成本从 O(N) 降到接近常数。业界高度关注这一方向,因为在大 batch 生产环境中,注意力成本已超过 FFN 成为主要瓶颈,尤其 agentic 系统需要反复读取长历史,若不优化,token 生成延迟和费用将难以承受。
行业类比
类似于在超长会议记录问答中,助手应能直接跳转到与问题相关的段落,而不是逐字浏览全部记录后再作答。
核心洞察
- Declarative Attention 将稀疏注意力的选择权从外部代理评分转移到模型自身,通过解析模型在链式思考中输出的控制标记(<global>/<focus>/<local>)直接跳过大部分 KV cache 读取。现有动态稀疏注意力方法依赖轻量代理分数预选 token,每步仍需 O(N) 扫描上下文计算分数,没有消除全局依赖;DA 则是模型内在声明相关区域,推理引擎像处理工具调用一样执行,零样本即可实现显著 token 节省,且不改变模型权重。
- 零样本 DA 的精度损失随模型规模缩小,表明强模型已具备自主声明注意力的能力,稀疏注意力无需外部训练即可自然涌现。实验显示 31B 模型有约 1.27pp 精度下降,但更大模型能更准确声明相关区域,与全注意力差距缩小,而 token 节省比例基本不随规模变化。这一观察为未来通过后训练或强化学习进一步增强声明能力提供了基础,也说明模型自身的上下文理解足以支撑声明式稀疏注意力框架。
方法
输入与提示设计
Declarative Attention (DA) 将长上下文分割为多个带标识的 segment,并在 prompt 中告知模型可引用的段编号。模型在 chain-of-thought 中生成特殊声明标签:<global>(全量上下文)、<focus:segment_id>(指定段)、<local>(仅最近输出)。
解码时干预
推理引擎把声明解析为类似 tool call 的控制信号,动态切换注意力 mask。mask 采用块对齐方式,仅对 global attention layers 施加,其余层保持全量 attention 以保证质量。集成到 vLLM 推理框架,通过自定义 kernel 跳过大部分 KV cache 读取:
<global>:读取全部 KV cache<focus:n>:仅读取第 n 个 segment 对应的 KV cache 块<local>:仅读取最近输出窗口对应的 KV cache
输出与效果
零样本下,Gemma-4-31B 和 Qwen-3.6-27B 在 15 个长上下文任务中 attended tokens 分别降低 52.0% 和 31.1%,准确率平均下降 1.27pp 和 2.75pp,且随模型规模增大差距缩小。
与预先用轻量 proxy score 选择 token 的外在稀疏方法不同,DA 是一种内在的、模型自主声明的注意力稀疏机制,无需额外训练或评分网络。
实验
实验设计
- 在 15 个长上下文任务上零样本评估 Declarative Attention。
- 使用 Gemma-4-31B 与 Qwen-3.6-27B 两个现成模型。
- 基线包括 vanilla global attention 与部分基于轻量代理分数的 extrinsic scoring 方法。
- 推理引擎使用定制 vLLM 集成,支持
<global>、<focus>、<local>三种模式解析。
关键发现
- DA 显著减少解码期间总 attended tokens:Gemma-4-31B 减少 52.0%,Qwen-3.6-27B 减少 31.1%。
- 对应精度下降分别为 1.27pp 和 2.75pp,且随模型规模增大而缩小。
- 注意力成本下降主要来源于 attention mask 本身,而非其他计算。
- 上下文长度扩展时,精度保持而 token 节省增大。
与基线对比
- 与 vanilla global attention 相比,DA 以极小精度代价换取约半数 KV cache 读取减少。
- 与 extrinsic proxy scoring 方法相比,DA 不引入每步 O(N) 的额外打分开销,而是依赖模型自身声明相关区域,更接近内在稀疏注意力。
- 结果支持 DA 作为 system-2 型稀疏注意力机制,为训练后优化及 agentic 系统提供新方向。
行业影响
落地场景
Declarative Attention (DA) 可直接用于长上下文推理服务,如 电商客服机器人 处理跨会话的多轮对话、企业知识库问答 扫描数万页文档、代码助手 阅读大型仓库历史。模型通过 <global>/<focus>/<local> 声明只需关注特定段落,推理引擎跳过无关 KV cache 读取,特别适合 agent 任务中反复调用长上下文的场景。
商业价值
降低推理成本是主收益。在 1M token 上下文中,DA 在 Gemma-4-31B / Qwen-3.6-27B 上减少 decode 阶段 attended tokens 52.0% / 31.1%,准确率仅下降 1.27pp / 2.75pp,且随模型规模增大差距缩小。这意味着长上下文 API 服务可降低 GPU 内存带宽压力,提升并发吞吐,或在相同成本下支持更长上下文,增强产品竞争力。
与现有工作流接口
DA 以 协议形式 嵌入推理引擎,类似 function calling 解析模型输出的 <global>/<focus>/<local> 声明,动态构建 block-aligned attention mask。可与 vLLM 等现有 serving 框架集成,作为 KV cache 管理的上游控制信号。企业可在现有推理栈中增加一个轻量解析层,无需重新训练模型,零样本即可启用。后续可配合稀疏注意力、speculative decoding 进一步放大收益。
局限
- 零样本准确率下降:论文报告 **Gemma-4-31B** 和 **Qwen-3.6-27B** 在 15 个长上下文任务上分别有 1.27pp 和 2.75pp 的准确率下降,虽然随模型规模缩小,但下降仍存在。对于准确性敏感的场景,这种有损的稀疏注意力可能不可接受。DA 依赖模型正确声明关注区域,若声明错误则可能丢失关键信息。此外,目前仅在两个开源模型上验证,泛化到其他架构或专有模型尚不明确。
- 协议依赖与集成复杂度:DA 需要模型遵循特定的声明协议(如输出 `<global>`、`<focus>` 等控制 token),这对模型的指令遵循能力有要求。对于较弱或未对齐的模型,声明可能不可靠。推理引擎必须解析这些控制 token 并动态构建注意力掩码,增加了系统复杂度。论文虽提供了 vLLM 集成,但自定义内核和对不同注意力实现的适配仍需额外工程工作,可能阻碍快速部署。
- 与训练型稀疏方法的差距:DA 作为零样本方法,未经过专门训练优化,其注意力节省和准确率可能低于经过训练的动态稀疏注意力方法。论文也承认未来基于训练的方法可以进一步挖掘潜力。此外,DA 目前仅应用于全局注意力层,未覆盖所有层,且声明区域的粒度较粗(块对齐),限制了更细粒度的 token 级节省。与 KV cache 逐出或轻量级代理打分等方法相比,DA 的优势依赖于上下文足够长,短上下文场景收益有限。