Language Models Can Control Their Own Attention | Summary
02 Sep 2026 | Paper Review Sparse Attention Long-Context Language Models KV CacheContents
- Summary
- 2. Declarative Attention
- 3. When does DA pay off?
- 4. Experimental settings
- 5. Results
- 6. Analysis
- 8. Discussion, Limitations, and Future Work
- Appendix
- Brief Thoughts
This article explains the key points of Language Models Can Control Their Own Attention.
- 2026-09-02 (arXiv)
- Ho, Namgyu, Ahmad, Huzama, Koh, Woosung, Yun, Se-Young, Schuster, Tal, Santos, Cicero Nogueira dos.
- KAIST AI, Google DeepMind
- Paper
Summary
- Declarative Attention (DA) is a zero-shot inference protocol for long-context language models that turns model-generated attention declarations into decode-time KV-cache masks. The model alternates between <global>, <focus>, and <local> spans; across 15 tasks, DA reduces attended tokens with modest average accuracy loss on sufficiently capable models.
2. Declarative Attention
DA makes attention scope explicit in the generated reasoning trace rather than predicting relevant tokens from proxy scores or hidden activations. A state machine parses declarations and changes visible KV-cache blocks while retaining the system scaffold, question, instructions, and generated response.
2.1. Prompt overview
The prompt separates an always-visible scaffold from the long context. The scaffold comprises a short system instruction, user question, and DA instructions, while the state machine selectively exposes the context region.
2.2. Context delivery
DA divides context into addressable magic chunks targeting 2048 tokens. The segmenter prefers paragraph, newline, sentence, clause, and then word boundaries, and presents chunks in a simulated tool-use transcript without executing retrieval tools.
2.3. Decode-time interventions
The runtime starts in global mode; opening <focus magic_chunks=”K”> or <local> changes the mask, and a matching closing tag returns it to global mode. Its vLLM integration rewrites the request KV-cache block table at each decode step, keeping whole blocks so existing attention kernels run unchanged.
The example traces an answer through global navigation, focused extraction from Magic Chunk 1, and local arithmetic. The mask reduces visible prompt content from all 25,466 prompt tokens to the scaffold plus one chunk, then to the scaffold and generated trace only.
3. When does DA pay off?
DA exchanges additional decoding steps for lower per-step global-attention memory reads. The paper argues that the trade-off is favorable in large-batch, long-context serving, where KV reads are memory-bound and grow with context length while FFN work is treated as compute-bound in the roofline analysis.
4. Experimental settings
Experiments cover six Gemma and Qwen models and 15 long-context tasks from RULER, LongBench, LooGLE, and ZeroSCROLLS. Contexts over 244K tokens, or 116K for Gemma-4-E4B, are excluded; each source contributes up to 128 fixed-seed examples, thinking is disabled, and generations are capped at 8K tokens.
The benchmark suite contains 15 sources divided into single-span retrieval/reasoning and multi-span reasoning. It spans contexts around 6–14K tokens through LBv2/code_repo, whose mean context length is 1071.1K tokens before input-length filtering.
5. Results
Results compare Vanilla full causal attention, DA-no-mask (DAnm), which uses the DA prompt with full attention, and DA with dynamic masking. Accuracy is judged with generated acceptance rubrics and Qwen-3.5-4B, while attended tokens sum KV positions accessed during decoding and roofline estimates target optimized B200 serving.
5.1. Main results
Across the 15 tasks, DA lowers Gemma-4-31B accuracy from 87.01% to 85.74% and Qwen-3.6-27B accuracy from 85.31% to 82.56%, while reducing attended tokens from 13.43M to 6.45M and from 22.54M to 15.52M per sample, respectively. DAnm nearly preserves accuracy but raises attended-token totals to 22.31M and 29.02M, showing that masking provides the net efficiency gain and most of DA’s accuracy cost.
The detailed task results show that DA lowers attended tokens relative to Vanilla on most sources, although task-level accuracy and savings vary. DAnm can substantially inflate attended-token totals despite near-lossless formatting, isolating the effect of masking.
Relative to vanilla, DA uses 135% of decode steps on Gemma-4-31B and 131% on Qwen-3.6-27B, yet attends to only 48% and 69% as many tokens. DAnm also adds steps but attends to 166% and 129% of vanilla tokens, identifying the mask as the source of savings.
5.2. Model capability scaling
Relative DA accuracy rises with model scale, from 29% to 99% of vanilla within Gemma 4 and from 64% to 97% within Qwen 3.5/3.6. Five of six models attend fewer tokens than vanilla; Gemma-4-12B is the exception because roughly 6% of DA responses exhaust the 8K budget, and excluding non-terminating responses lowers it to 98% of vanilla attended tokens.
Relative accuracy approaches vanilla as models grow in both families, while attended-token reductions are broadly comparable across scales except where non-termination inflates totals. The Gemma-4-12B anomaly is largely removed when non-terminating responses are excluded.
5.3. Context-length scaling
On Gemma-4-31B, DA stays within about one percentage point of vanilla through 32K-token contexts, then declines to roughly 96% relative accuracy in the 64–256K bin. Its absolute saving grows from about 1M attended tokens at short contexts to about 21M in the longest bin because masking retains a broadly constant fractional reduction as context grows.
Pooling Gemma-4-31B examples by context length shows a modest long-context accuracy divergence for DA that is absent from DAnm. Meanwhile, DA's attended-token advantage increases in absolute magnitude to about 21M fewer tokens in the longest bin.
5.4. Efficiency results
Using a roofline estimate for one B200 accelerator in bf16, with MFU 40% for matrix multiplies and MBU 70% for memory reads, DA reduces estimated decode time from 269.1 ms to 192.3 ms on Gemma-4-31B and from 306.2 ms to 237.3 ms on Qwen-3.6-27B. These are theoretical decode-only estimates, not measured end-to-end latency, and assume optimized disaggregated large-batch serving.
The decomposition attributes most vanilla decode time to global-memory reads: 196.5 of 269.1 ms for Gemma and 263.8 of 306.2 ms for Qwen. DA lowers that term while its added decoding increases matmul and local-memory costs, producing estimated totals of 192.3 ms and 237.3 ms.
6. Analysis
The analysis attributes DA’s savings to time spent in focus and local modes and measures whether models generate valid focus declarations. It distinguishes zero-shot protocol limitations, including verbosity and invalid references, from the block-masking mechanism.
6.1. Efficiency impact of DA modes
On Gemma-4-31B, about 27% of generated tokens are global, while focus and local account for the other 73%. Focus and local attend to about 12% and 6% of vanilla-step tokens on average, yielding 76–99% per-token savings; global use nevertheless rises to about 45% in the longest context bucket.
Focus and local modes dominate generation and sharply reduce per-token reads, explaining savings despite longer outputs. The rise in global-mode share at longer contexts limits aggregate savings.
6.2. Protocol adherence
Focus-call validity rises with model size, from 58% for Gemma-4-E4B to 99% for Gemma-4-31B and from 89% for Qwen-3.5-4B to 99% for Qwen-3.6-27B. Focus attempts remain around 1.4–1.9 per response, indicating that larger models chiefly improve by emitting valid chunk references rather than by issuing fewer focus calls.
Valid focus references become nearly universal for the largest models, while the number of attempts does not consistently decrease with scale. This separates protocol-adherence improvements from simply avoiding focused retrieval.
8. Discussion, Limitations, and Future Work
The paper presents DA as a reversible and explicit alternative to latent sparse-attention selection: the full cache remains resident, but the model controls which blocks are read. Limitations include residual global-mode cost, zero-shot decomposition quality, artificial benchmark segmentation, non-thinking-mode evaluation, and roofline projections rather than measured end-to-end latency.
8.2. Further potential of DA with post-training and in agentic settings
The authors propose post-training DA policies to reduce the approximately one-third increase in decoding steps and better fit task structure. They suggest that agentic systems offer natural addressable units such as tool outputs, user turns, and retrieved passages, but do not evaluate these settings.
8.3. Synergy with modern efficient inference techniques
DA can compose with lightweight scan-based sparse attention by applying scans during expensive global spans while using declarations during focus and local spans. It may also complement speculative decoding because masking lowers KV reads per verification pass while speculation reduces sequential decode passes.
Appendix
- The appendix details related sparse-attention and KV-cache methods, the vLLM implementation, roofline derivations, datasets, judge validation, sampling, additional results, and full prompts. It identifies two structural failure classes: segmentation can destroy necessary evidence, and document-proportional outputs can eliminate end-to-end savings despite lower per-step attention.
Brief Thoughts
DA provides an interpretable sparse-decoding mechanism because the model’s textual plan directly controls the attention mask. The strongest evidence concerns attended-token reductions on capable models; practical deployment claims still require end-to-end latency measurements, thinking-mode evaluation, and improved policies for global navigation and structurally difficult tasks.