Back Home

推論系統

Declarative Attention Lets Models Specify KV Cache Read Ranges, Cutting Tokens Attended During Long-Context Decoding by Over 30%

KAIST and Google DeepMind have introduced Declarative Attention, which allows existing models to switch among global, focused, and local attention modes within their reasoning text. The authors reduced attended tokens by 31.1% to 52.0% across two models, though the speed gains remain model-based estimates and accuracy declined slightly.

Segnargsed · CC BY-SA 4.0 · Image source
zh-Hant

During long-context model decoding, generating each token typically requires reading a massive KV cache from GPU memory. Even when an answer involves only a few passages, the lookup cost still grows with context length. Researchers from KAIST and Google DeepMind have proposed Declarative Attention (DA), which turns “where to look” into a control protocol that the model itself can output, without introducing a separate retriever that must scan the entire context.

The system first divides the input into addressable blocks of roughly 2,000 tokens. During reasoning, the model emits `<global>`, `<focus>`, or `<local>` tags. Global mode can access all blocks; focus mode reads only specified blocks; and local mode does not read the original long text, retaining only the question, fixed instructions, and previously generated content. The inference engine parses these tags with a state machine and dynamically rewrites the KV-cache block table, allowing existing FlashAttention-style kernels to skip blocks that do not need to be loaded. The method requires no model fine-tuning and avoids the cost of an external selector performing an O(N) scan at every step.

The team conducted zero-shot evaluations of Gemma-4-31B and Qwen-3.6-27B on 15 long-context tasks. Average attended tokens fell by 52.0% and 31.1%, respectively, while accuracy declined by 1.27 and 2.75 percentage points. Roofline analysis estimated that decoding time could fall to 0.71× and 0.77× the baseline, respectively, but these figures are not measured end-to-end throughput for a live serving system.

The key engineering question is whether models can reliably and honestly indicate their information needs. Smaller models were markedly less capable of following the protocol; global steps may still dominate total cost, and selecting the wrong blocks can directly hide relevant evidence. The paper currently provides no public implementation. Future work should examine real-world GPU latency, compatibility with batched serving, and whether training models to use the attention tags can narrow the accuracy gap.

Sources

  1. Language Models Can Control Their Own Attention
  2. Large Language Models Can Control Their Own Attention Span