来自谷歌DeepMind及其合作者的重要论文。（务必收藏）该模型在每个生成的token上都会读取其整个KV缓存，尽管最终只会关注其中极小的一部分。换言之，即使你仅询问一段百万token对话中的某个细节，全局注意力层在生成每个token时都必须重新阅读全部内容。常规的解决方案是先用廉价的代理评分来猜测相关token，但这每一步仍需O(N)的开销。声明式注意力则改用询问模型的方式：模型在自身的思维链中声明需要查看的位置。这样，生成过程被划分为三种模式：全局模式读取完整上下文，聚焦模式读取特定区域，局部模式仅读取最近的输出。推理引擎像解析工具调用一样解析这些声明，并跳过大部分缓存读取。在针对现成权重模型的15项长上下文任务零样本测试中，解码时的被关注token数量在Gemma-4-31B上降低了52.0%，在Qwen-3.6-27B上降低了31.1%。

论文地址：https://arxiv.org/abs/2609.02737 论文讨论：https://academy.dair.ai/papers/language-models-can-control-their-own-attention-2609.02737

语言模型能够控制自身注意力

来源：https://arxiv.org/html/2609.02737 报告编号：Huzama Ahmad (隶属：KAIST AI) Woosung Koh (隶属：KAIST AI) Se-Young Yun (隶属：KAIST AI) Tal Schuster (隶属：Google DeepMind) Cicero Nogueira dos Santos (邮箱：[email protected] 通讯作者) 隶属：Google DeepMind

摘要

摘要：语言模型将其大部分注意力集中在上下文的一小部分上，但它必须读取整个KV缓存才能找到那几个关键的token。如果用户询问一段百万token对话中的某个历史细节，全局注意力层必须扫描完整上下文以生成回复的每一个token。一种主流方法通过轻量级代理评分预先选择相关token来缓解这种成本，但这种外部评分方法每一步仍然需要O(N)的复杂度。我们采取了一种内部方法，其动机源于一个简单的问题：模型难道不是已经知道上下文的哪些部分是相关的吗？为此，我们引入了声明式注意力，一种引导模型在其思维链中声明所需关注位置的协议，将生成过程划分为三种模式：全上下文、特定区域和仅最近输出。推理引擎像解析工具调用一样解析这些声明，并跳过大部分KV缓存读取。在针对15项长上下文任务的零样本评估中，在现成模型上应用声明式注意力显著减少了解码期间的总关注token数量（Gemma-4-31B减少52.0%，Qwen-3.6-27B减少31.1%），同时准确率下降幅度温和（分别为1.27和2.75个百分点），且该差距随模型规模扩大而缩小。声明式注意力开辟了稀疏注意力的新维度，未来工作通过基于训练的方法有望进一步挖掘其潜力。

参见图注 图1：声明式注意力让模型能够控制自身的注意力范围。 (提示)一段系统指令、分为“魔法块”的长上下文信息、用户问题以及声明式注意力指令。 (模型回复)模型自由地在三种模式间切换：浏览完整上下文、在指定块上推理、在回复中已累积的信息上推理。我们的协议不限制每种模式的使用方式、时机或次数。 (声明式注意力掩码)根据模型回复文本和声明式注意力语法，推导出段级注意力掩码。每一行分别代表全局、聚焦和局部推理步骤。支撑结构（系统、问题、指令）在每种模式下均保持被关注，而上下文段在除该模式保留的段之外均被遮蔽（虚线）：全局模式下保留所有段，聚焦模式下仅保留标记为F的聚焦段，局部模式下不保留任何段。声明式注意力无需辅助评分器，可在无需训练的现成模型上运行。

1 引言

Transformer在每个解码步骤都对每个前置token计算注意力，这使其在长上下文任务中计算成本高昂。具体而言，键值缓存的内存访问延迟在长上下文场景中严重主导了解码时间。例如，在1M token上下文中，Qwen-3.5-397B-A17B每个解码步骤大约需要加载15GB的KV缓存；其内存带宽需求与加载模型的17B活跃参数相当。然而，这种穷尽机制与人类认知存在根本分歧：当理解一篇长文档时，人类不会持续重读每一个先前的单词来综合信息并回答问题。实证研究证实了这一观点，表明注意力分数自然集中于一小部分上下文token。虽然将上下文注意力稀疏化是一个自然的解决方案，但一个根本挑战依然存在：真正的注意力分数是先验未知的。因为这些分数只有在计算完整注意力矩阵后才可知，动态识别需要关注哪些token的成本高得令人望而却步。先前的工作试图通过预测高注意力token并应用掩码来绕过这一问题，以排除无关上下文的计算和内存访问。早期方法依赖静态启发式方法，如时间近邻性或历史注意力幅度。然而，这些固定规则难以预见未来查询所需的具体token，导致长上下文性能下降。相反，最近的方法承认了注意力的查询依赖性，在每个解码步骤通过轻量级扫描KV缓存来近似掩码。虽然这降低了成本的常数因子，但每一步的复杂度仍然是O(N)。我们将在附录A中全面调查此及相关工作。

贡献

我们提出一个正交的方向：引导模型明确声明它将在每一步关注何处。最近的研究表明，语言模型在其隐藏状态中编码了关于未来token的信息，而思维链提示将这种潜在计算呈现为可解释的文本。我们扩展了这一原则，从指定思考什么扩展到关注何处。令人鼓舞的是，先前的工作表明，这种行为可以通过在2K token上下文中针对手工设计的上下文分区和注释语法分别进行微调来训练进模型。我们证明，现在可以通过单一的任务无关协议来引导这种行为：现代现成模型在固定提示下能够零样本地声明其注意力，适用于广泛的、上下文长度达约100K token的长上下文任务。通过直接从模型生成的推理轨迹推导注意力掩码，而不是通过隐藏激活近似，我们的方法完全消除了选择成本。O(N)的上下文读取仅存在于模型声明的全局阶段，而非每一步的开销。具体而言，我们引入了声明式注意力协议，它将生成划分为三种不同的注意力模式：(1) 全局模式，模型浏览完整上下文；(2) 聚焦模式，模型专注于特定上下文区域；(3) 局部模式，模型仅关注自身最近的输出，而不查看完整上下文。模式转换以可解析的token形式在思维链中发出，推理引擎读取这些token以在每个解码步骤动态构建注意力掩码。

我们通过实验展示了声明式注意力的价值（第4节，第5节）：

1. 零样本有效性作为下限：声明式注意力在现成模型上零样本工作，无需参数更新。因此，我们的结果代表了性能的下限，如果模型针对该协议进行后训练，预期有显著提升空间（第8节）。
2. 有利的成本-准确率权衡：在15项长上下文任务中，声明式注意力将平均解码注意力成本在Gemma-4-31B上降低了52.0%，在Qwen-3.6-27B上降低了31.1%，仅导致轻微的准确率下降（分别为1.27和2.75个百分点）。
3. 正向的缩放与成本节约：声明式注意力直接受益于模型能力，随着规模从4B增加到31B，准确率差距稳步缩小。此外，随着上下文增长，绝对token节约量急剧增加（每次回复最多节省21M token），消融研究证实，是动态掩码本身（与无掩码消融相比，最多削减71.1%的关注token）驱动了节约，而非提示格式。
4. 高效的vLLM实现：我们将声明式注意力集成到vLLM中，采用了与FlashAttention兼容的块对齐、原地KV缓存遮蔽。基于性能上限模型的墙钟时间分析预测，在优化良好的服务栈上，声明式注意力的注意力节约将使解码墙钟成本降至原始方案的0.71倍（Gemma-4-31B）和0.77倍（Qwen-3.6-27B）。

2 声明式注意力

在普通大语言模型前向传播中，大部分注意力权重集中在上下文的一小部分上，并且这种稀疏模式随着步骤变化，其方式无法提前预知。因此，每个解码步骤都必须从HBM读取整个KV缓存，即使其中大部分对输出的影响微乎其微。声明式注意力协议引导模型重构其思维链，使其注意力计划变得可读。它要求模型：（i）将其推理组织成注意力范围保持稳定的连续片段，并（ii）使用预定义的标签语法声明该范围。为实现相应的掩码，我们引入了一个与推理引擎并行的声明式注意力状态机：它监视模型输出中的标签转换，并在每个解码步骤更新注意力掩码。

声明式注意力定义了三种模式：、和，每种服务于不同的推理目的。对问题“Acme公司成立多久后上市？”的回答可能如下： 我需要成立年份和上市年份。魔法块2中的公司历史应该陈述成立情况。 “Acme公司于2003年在圣何塞成立。” 上市年份仍然缺失。魔法块7涵盖了Acme的财务里程碑。 “Acme公司于2011年在纽约证券交易所上市。” 2011 - 2003 = 8年。8年。 聚焦声明中命名的魔法块是长输入的可寻址段，我们通过将上下文分割成约2K token的单元来准备这些段。在每种模式下，模型都持续关注问题、指令以及到目前为止自身的回复。模式之间的区别仅在于上下文（提示的分段长输入区域）的可见范围：

• 关注所有上下文段。这是导航模式，用于浏览完整上下文以定位相关信息。我们的零样本提示指示模型使用此模式识别下一个要聚焦的段，并简要记录该段相关的原因。
• 仅关注标签中命名的上下文段。这是对特定上下文区域进行推理而不承担关注其余部分成本的模式。我们的提示指示模型从命名段中逐字提取所需值。
• 不关注任何上下文段。这是对回复中已有信息进行自包含推理的模式。我们的提示指示模型针对问题进行规划，并根据先前提取的值综合最终答案。

声明式注意力本身不限制每种模式的使用方式、时机或次数。由于我们从现成模型零样本引导声明式注意力，我们的提示为每种模式提供了指导以支撑其声明式推理。本节的其余部分描述：（1）提示结构，（2）上下文如何作为可寻址段传递给聚焦模式，以及（3）声明式注意力状态机在解码时对推理引擎的干预。

2.1 提示概述

我们围绕两个区域设计提示：一个支撑结构始终对模型可见，以及一个上下文包含模型必须推理的长输入。支撑结构在每种模式下均保持关注，为声明式注意力协议提供持久的基准。上下文是其可见性由声明式注意力状态机控制的可变区域，在长上下文任务中包含提示的绝大部分token。图1展示了提示结构，包括围绕上下文的三个支撑部分（标记为蓝色）：

• 系统指令：一个简短的固定序言，其内容填充注意力槽，使上下文永远不会进入其中。
• 上下文：模型推理所用的长输入，以模拟工具使用记录中的可寻址段形式传递。
• 问题：用户的查询。
• 指令：定义三种模式，并根据提示开发期间观察到的失败模式提供使用指南。

完整的声明式注意力指令提示及其普通对应物普通指令提示见附录F。

2.2 上下文传递

要使用聚焦模式，模型必须命名其打算关注的上下文区域。因此，我们将上下文划分为可寻址的段。寻址一段需要模型跟踪该段的起始和结束位置。大型语言模型训练所用的文本具有普遍的结构：预训练语料库包含文档和段落边界，后训练对话包含用户、助手和工具轮次边界。我们期望模型能够追踪这些边界以进行段寻址。