Letting Language Models Decide Where to Look

Introduction

Long-context inference is constrained not only by how many tokens a model can store, but also by how much of that context it must read for every generated token. A standard global-attention layer may scan the entire KV cache even when the answer depends on only a small region. In a conversation approaching one million tokens, repeatedly reading that cache can become a major systems cost.

A team from KAIST AI proposes Declarative Attention, or DA, as a way for the model to state where it needs to look. Rather than adding a separate retrieval stage, DA defines a protocol through which the model and inference engine coordinate context access.

How it works

DA separates generation into three modes:

• global: read the full context when broad access is necessary;
• focus: attend to a specified region of the context;
• local: read only recent output when earlier material is not needed.

During chain-of-thought generation, the model emits declarations corresponding to these modes. The inference engine parses them in a manner similar to tool calls and uses them to decide which portions of the KV cache to read. If the next reasoning step is local, the engine can avoid scanning the full history.

This differs from methods that first assign lightweight proxy scores to all candidate tokens and then select a subset. Such scoring can reduce the eventual read volume, but it may still require processing a context of length N at every step. DA instead asks whether the model already has enough information to identify the relevant region and can therefore make the selection part of generation itself.

Results and caveats

The supplied material describes a zero-shot evaluation across 15 long-context tasks using off-the-shelf Gemma-4-31B and Qwen-3.6-27B models. It reports reductions in total attended tokens of 52.0% and 31.1%, respectively. The accompanying accuracy costs are described as modest: 1.27 percentage points for Gemma and 2.75 points for Qwen in the paper summary. However, a separate page snippet supplied with the material lists a 1.52-point drop for Gemma, so that figure should be checked against the final paper version.

The trade-off is important. A wrong declaration can direct the engine away from information needed for the answer. The material also says that the accuracy gap shrinks with model scale, but the evaluation is zero-shot and does not establish how reliably models can learn this protocol with dedicated training.

Why it matters

DA introduces a distinct route to sparse attention: the model itself participates in deciding its context-access pattern. That could be useful for million-token conversations, long-document analysis, and workloads where output generation is long enough for repeated KV-cache reads to dominate.

There are still practical questions. A production system would need robust declaration parsing, safeguards for incorrect focus regions, and a recovery strategy when the model realizes it has looked in the wrong place. It must also show that fewer attended tokens translate into lower end-to-end latency and cost, rather than merely moving work elsewhere.

For now, Declarative Attention is best understood as a research demonstration of model-controlled context access, not a universal solution to long-context inference. Its most consequential contribution may be conceptual: sparse attention need not be imposed entirely from outside; a capable model may be able to request the context it needs.

Source: Hugging Face Daily Papers