Paper Walkthrough: Language Models Can Control Their Own Attention

During long-context decoding, a model re-reads its entire KV cache at every step. Declarative Attention has the model announce, inside its own chain-of-thought, where it will look next, and the inference engine builds an attention mask from that text. A walkthrough of the paper that cut attended tokens by 52.0% and 31.1% on off-the-shelf models, zero-shot.

Language Models Can Control Their Own Attention

The librarian who re-reads 10,000 pages for every word

Picture 10,000 pages stacked on a desk. A reader asks, "How many years after founding did this company go public?" The librarian flips through every page from first to last, says "eight" — and then flips through all of them again just to say the next word, "years."

That is roughly what a language model does while decoding (generating one token at a time) over a long context. The paper we're reading is titled "Language Models Can Control Their Own Attention" (arXiv:2609.02737, KAIST AI and Google DeepMind, published 2026-09-02).

Here is its claim in brief. Language models spend most of their attention on a small fraction of context, yet they read the entire KV cache just to find the few tokens that matter. A prominent mitigation pre-selects relevant tokens using lightweight proxy scores, but that extrinsic scoring itself costs $O(N)$ per step. So why not go intrinsic — wouldn't the model already know which parts of the context are relevant? The authors introduce Declarative Attention (DA), a protocol that elicits the model to declare where it needs to attend within its chain-of-thought, partitioning generation into three modes: <global> (full context), <focus> (a specific region), and <local> (recent output only). The inference engine parses these declarations like tool calls and skips most of the KV cache read. Under zero-shot evaluation across 15 long-context tasks, DA on off-the-shelf models (Gemma-4-31B, Qwen-3.6-27B) reduces total attended tokens during decoding by 52.0% and 31.1%, with modest accuracy drops of 1.27pp and 2.75pp — drops that shrink as the model scales.

Why "read everything" is the default

In the long-context regime, the latency of reading the Key-Value (KV) cache dominates generation time. The paper's example: on Qwen-3.5-397B-A17B with a 1M-token context, roughly 15 GB of KV cache must be loaded per sequence at every decoding step — a memory-bandwidth requirement comparable to loading the model's 17B active parameters. Attention scores are empirically shown to concentrate on a small subset of context tokens, so why read all of it? Because the true attention scores are unknown a priori: they only become available after the full attention matrix has been computed (§1). If the KV cache is new to you, see understanding the KV cache from scratch.

Prior workarounds, and where they stop

Earlier work predicted which tokens would attract attention and masked out the rest. The first generation used static heuristics — recency, or historical attention magnitude — which cannot anticipate the tokens a future query will need, so long-context performance degrades. Newer methods approximate the mask with a lightweight scan over the whole KV cache at each decode step. That lowers the constant factor, but the per-step complexity is unchanged (§1):

Equation (1) says that the work of deciding where to look scales directly with the number of context tokens $N$. At a million tokens, you touch a million tokens' worth of state every single step just to make the decision. That is the paper's starting point.

The three modes: let the model declare

Language models encode information about future tokens in their hidden states, and chain-of-thought prompting surfaces that latent computation as readable text. DA extends the principle from dictating what to think to dictating where to attend (§1).

<global> attends to all context segments — the navigation mode, used to identify the next segment worth focusing on and briefly note why. <focus> attends only to the segments named in the tag — the close-reading mode, used to extract needed values verbatim. <local> attends to no context segments at all — the self-contained reasoning mode, used to synthesize the final answer from values already extracted into the response. In every mode the question, the instruction, and the model's own response stay attended; the only thing that changes is how much of the long input remains visible (§2).

<global>
I need the founding year and the IPO year. The company history in Magic Chunk 2 should state the founding.
</global>

<focus magic_chunks="2"> "Acme Corp was founded in 2003 in San Jose." </focus>

<local> 2011 - 2003 = 8 years. </local>

<answer> 8 years </answer>

You can read it. That is the point: the tokens that determine the mask are the same tokens a human can audit. DA needs no auxiliary scorer and runs on off-the-shelf models with no training (Figure 1).