Language Models Can Control Their Own Attention — part 3/6

Cleaned-text code points [30954, 46015). Bibliography omitted here; source.md includes it.

DA lowers estimated decode wall-clock time to 0.71 × \times of vanilla on Gemma-4-31B and 0.77 × \times on Qwen-3.6-27B at the stated B200 utilizations. The saving comes entirely from the global-memory read, which is the dominant decode cost under vanilla, 73% of decode time on Gemma and 86% on Qwen, and the only term the mask reduces. The matmul and local-memory costs are context-independent and in fact rise slightly under DA, because the DA protocol generates more decode steps than vanilla, 31% more on Qwen and 35% more on Gemma. The two models differ in how much of the global saving survives. On Qwen the local read is the GDN state, a small fixed cost at 5% of DA’s attention time, so the global reduction passes through almost undiluted. On Gemma the local read is the SWA floor, which spans 50 of the 60 layers and is 42% of DA’s attention time, large enough to cap the saving even though the global read itself falls by more than half. These figures are a theoretical ceiling at the stated utilizations rather than measured wall-clock time, and they exclude prefill, which under phase-disaggregated serving runs on a separate pool. Section C derives every entry.

## 6 Analysis

### 6.1 Efficiency impact of DA modes

As shown in Section 5 , <global> accounts for only about 27% of generated tokens on Gemma-4-31B , while <focus> and <local> together account for the remaining 73%. These are the cheap modes: <focus> and <local> attend to only a small fraction of the tokens a vanilla step attends (about 12% and 6% on average), yielding roughly 76–99% per-token attention savings. The <global> share does rise with context length, to about 45% in the longest bucket ( Section 5(a) ), which caps the total saving there. As with the mode mix in Section 3 , this reflects how the model reacts to the zero-shot prompt rather than a limit of the mechanism ( Section 8 ).

As expected, <local> , which attends only to the model’s generated tokens, achieves the largest savings, closely followed by <focus> . The savings of both modes increase with context length, with especially sharp gains for <focus> . These results suggest that DA’s efficiency benefit grows as long-context and long-horizon tasks continue to scale.

### 6.2 Protocol adherence

DA depends on the model emitting parseable control tokens, so we measure how reliably each model follows the protocol ( Section 6 ). We report the focus success rate, the fraction of <focus> calls that resolve to a valid chunk reference, and the number of focus attempts per response. Focus success rate rises with model size in both families, from 58% at Gemma-4-E4B to 99% at Gemma-4-31B and from 89% at Qwen-3.5-4B to 99% at Qwen-3.6-27B . The largest models follow the protocol almost perfectly, and adherence is the dominant bottleneck only at the smallest scale: Gemma-4-E4B ’s low success rate mirrors its accuracy collapse in Section 3 , confirming that much of that collapse is a failure to follow the protocol rather than a failure to reason. Focus attempts per response stay in a narrow band (about 1.4 to 1.9) with no clear trend in model size, indicating that stronger models succeed by parsing their focus calls more reliably, not by issuing fewer of them. The number of attempts is itself a zero-shot artifact, set by how the model chooses to use the protocol rather than by what is optimal for the task ( Section 8 ).

## 7 Conclusion

We introduced Declarative Attention, a zero-shot protocol that lets an off-the-shelf model declare where it will attend as part of its chain-of-thought, so the inference engine reads an attention mask off the model’s own text rather than approximating one from its activations. Across 15 long-context tasks, this mask supplies essentially all of DA’s attention savings at an accuracy close to vanilla, and it lands on the global-attention KV read that dominates decode time as context grows. Because we elicit the behavior entirely through prompting, these results are a floor rather than a ceiling ( Section 8 ): the gap to vanilla narrows as the backbone scales, and the token saving widens with context. This shows that a model can control its own attention span in text a person can also read, turning selective attention from a pattern inferred inside the network into one the model states out loud, enabling long-context inference that is at once legible and more efficient.

## 8 Discussion, Limitations, and Future Work

### 8.1 Relevance of DA in the present and future

DA rests on the premise that a model can reason over a long context while reading only a small part of it at each reasoning step. Our results show that current off-the-shelf models can adopt such a strategy simply through zero-shot prompting. Most decode steps run in focus or local mode, reading a single ∼ {\sim} 2K-token segment or only the response so far, saving up to 98% and 99% of the per-step attention read at contexts beyond 128K ( Section 5(b) ). DA’s relevance therefore reduces to whether global attention is a major factor in decode cost at all. On the models we evaluate, reading the entire KV cache for vanilla global attention at every decode step is the dominant decode cost, at 73% and 86% of roofline-estimated decode time on our benchmark traces with contexts up to 244K tokens ( Section 5.4 ). Below, we ask whether this remains true across current and future architectures and in agentic workloads.

#### Relevance of global attention costs in current and future architectures.

Architectures continue to optimize the size and access cost of the KV cache. Recent advances replace global attention in most layers with local and linear attention (SWA, GDN) ( Gemma Team et al., 2026 ; Yang et al., 2025b ) , shrink each entry through low-rank representations (MLA) and token merging ( DeepSeek-AI, 2024a ; DeepSeek-AI, 2026 ) , and reduce the read itself through lightweight indexers ( DeepSeek-AI, 2025b ) , together cutting the KV bytes read per context token by up to two orders of magnitude ( Section C.9 ). To identify the extent of these efficiency gains, we calculate the relative cost of global attention under our roofline model on very recent architectures. We find that attention accounts for 94% of decode wall-time at 1M-token contexts for Kimi-K3 ( Kimi Team, 2026 ) , the one model that retains vanilla global attention, and 56–97% for the indexer-based sparse designs ( DeepSeek-AI, 2026 ; Z.ai, 2026b ; MiniMax, 2026 , DeepSeek-V4-Pro, DeepSeek-V4-Flash, GLM-5.3-Flash, and MiniMax-M3;) ( Section C.10 ). 1 1 1 Smaller models geared toward on-device deployment, such as the ones studied in our work, are less aggressive with global attention optimizations, but we believe these architectures are more representative of our target setting of large-scale long-context serving. We also note that context windows have historically lengthened as per-token costs fell, and 1M-token windows are now offered at the frontier. If this trend continues, we believe aggregate attention costs will remain significant even as per-token costs drop.

#### Relevance of DA in agentic systems.

In agentic systems, context is retrieved incrementally through tool calls rather than provided as one long input, which may suggest that retrieval already solves the problem DA addresses. Retrieval, however, addresses a different problem: it decides what enters the context, while DA decides what the model attends to among what has already entered. If anything, retrieval compounds the cost that DA targets: each tool result remains in context for the rest of the episode, while it is typically relevant only to the steps that requested it. As these results accumulate, the context grows long and carries forward across turns. Therefore, we believe agentic contexts fall in the regime where DA’s savings are largest.

### 8.2 Further potential of DA with post-training and in agentic settings

#### Zero-shot DA already demonstrates significant attention savings and positive scaling.

All results in this paper come from prompting alone, on standard long-context benchmarks with fixed contexts and thinking disabled. Under this zero-shot protocol, DA already reduces attended tokens per response by 52% and 31% on our largest Gemma and Qwen models at accuracy close to vanilla ( Section 2 ). These gains improve along the two axes the field is already scaling: accuracy converges toward vanilla as the backbone grows ( Section 5.2 ), and absolute savings grow with context length ( Section 5.3 ). The remaining limitations are properties of this zero-shot elicitation rather than of the masking mechanism. We outline these limitations and the future directions we believe can address them below: • Suboptimal zero-shot strategies. DA runs about a third more decode steps than vanilla, and its decomposition does not always fit the shape of the task ( Section 2 ). Vanilla chain-of-thought reasoning has been heavily optimized through supervised fine-tuning and reinforcement learning ( OpenAI, 2024 ; DeepSeek-AI, 2025a ) . We believe a similar pipeline could elicit more optimal use of DA.

• Manufactured segments. Fixed-context benchmarks force us to split the context into artificial chunks ( Section 2.2 ). Agentic contexts instead provide natural, already-addressable segments such as tool calls, user turns, and retrieved passages ( Google, 2026 ; Anthropic, 2026c ; OpenAI, 2026a ) . DA may prove more natural in these settings.

• Non-thinking mode only. Models failed to follow the protocol within thinking traces in preliminary experiments, so all our results use non-thinking mode. However, current models are trained to reason between tool calls with interleaved thinking ( Anthropic, 2026b ; Moonshot AI, 2025 ) , so we believe models could readily use DA within thinking traces when its operations are exposed as standard tool declarations. This would extend DA’s savings to the regime where reasoning traces are longest.

• Global mode costs. Global steps retain full attention costs and account for over 80% of DA’s attended tokens, rising further with context length ( Section 5 ). We posit that navigation does not require a full-fidelity view of the context: the model could instead survey an in-context index, i.e., a short description of each segment. Since such an index can be orders of magnitude smaller than the context itself, we believe it can reduce these costs accordingly.

### 8.3 Synergy with modern efficient inference techniques

#### Lightweight scan sparse attention methods can alleviate costs in global mode.

Sparse attention methods based on a lightweight scan score a compact surrogate of every context token at each step and attend only to the highest-scoring subset ( Tang et al., 2024 ; DeepSeek-AI, 2025b ) . Such a scan could serve exactly the global steps, where DA retains full attention costs, while declarations carry the savings in the remaining modes. Integration should be light, as DA applies its mask by rewriting the KV-cache block table without touching the attention kernel ( Section 2.3 ).

#### Synergy with speculative decoding.

Speculative decoding and DA reduce complementary parts of decode latency: DA lowers the cost of each decode step by masking the KV read, while speculative decoding collapses many sequential steps into one verification pass, offsetting the extra steps DA runs ( Section 2 ). Each verification pass still reads only the masked KV, so the two savings stack. DA adds little complication: the attention pattern is fixed within a mode, so drafting proceeds under a single mask and speculation resets only at infrequent mode transitions.

### 8.4 DA as system-2 sparse attention

DA can be read as a system-2 form of sparse attention, in which the model decides what to attend to by reasoning about it explicitly in its own output. 2 2 2 Our system-2 framing follows System 2 Attention ( Weston and Sukhbaatar, 2023 ) , which treats attention selection as a task for the model’s explicit reasoning, though it realizes selection by regenerating the input once before answering rather than by masking attention as reasoning unfolds. Learned sparse attention methods, including the lightweight scans above, instead infer a keep-or-drop mask from internal activations at every step ( Zhang et al., 2023 ; Tang et al., 2024 ) . In this sense, they perform system-1 attention masking. Because DA’s selection is expressed in language, the strategy can be changed by instruction alone, and it improves with model scale without any change to the protocol ( Section 5.2 ). For the same reason, DA could be optimized like system-2 reasoning itself, with reinforcement learning that rewards both accuracy and attention efficiency ( OpenAI, 2024 ; DeepSeek-AI, 2025a ) , which we consider for future work. The attention plan is also legible by construction: the tokens that make it auditable are the same tokens that drive the KV reads, a property that serves oversight as much as efficiency ( Korbak et al., 2025 ) . Self-Selected Attention Span, which trains task-specific span selection ( Jin et al., 2024 ) , notes this interpretability benefit as well.

#### Toward system-2 KV cache offloading.

DA changes its attended set only at span boundaries, and each change is announced in plain text before the tokens that read it. This is the access pattern KV cache offloading needs: out-of-focus segments could spill to host memory for the duration of a span and be prefetched back as soon as a declaration names them, overlapping the transfer with decoding. Existing systems offer no such handle: serving-level offloading moves KV between memory tiers only at request boundaries ( Sheng et al., 2023 ; Qin et al., 2024 ) , while sparse attention methods reselect tokens at every step, leaving nothing to prefetch ahead of time ( Tang et al., 2024 ; Lee et al., 2024 ; Chen et al., 2025 ) . This capability further connects DA to context compaction, where agentic systems summarize or clear parts of the history to remove redundant content ( Anthropic, 2025 ; Anthropic, 2026a ; OpenAI, 2026b ) . Compaction is destructive: the summary replaces the original tokens, their KV cache is discarded, and bringing dropped content back requires re-prefilling it from text. Under DA, the sequence is never edited, so the cache of out-of-focus content stays valid even while it sits off-device. We believe this can enable reversible context compaction , where the model dynamically offloads and restores KV cache to remove redundancy, without the need for re-prefill.

## Acknowledgements

We thank Yun-hsuan Sung, Ziwei Ji, Ganesh Jawahar, Vincent Perot, Santiago Ontañón and Yujin Kim for extensive discussions. We also thank the Google Cloud Research Credits program for supporting this project.

This work was supported by the Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korean government(MSIT) (No. RS-2024-00457882, National AI Research Lab Project).

## Appendix A Related works

### A.1 Dynamic sparse attention at decode