AIAny
Icon for item

Language Models Can Control Their Own Attention

Introduces Declarative Attention (DA), a zero-shot protocol that has LMs declare which parts of long context to attend to during chain-of-thought, letting the runtime build dynamic attention masks and skip most KV-cache reads. Produces large token savings (up to ~52% on Gemma-4-31B) with modest accuracy loss.

Introduction

Long-context transformer decoding is dominated by KV-cache reads even though models concentrate attention on a tiny fraction of tokens. This paper flips the usual approach: instead of externally scoring relevance, it elicits off-the-shelf models to state where they will attend as part of their reasoning. The runtime parses those declarations into a dynamic attention mask, avoiding O(N) token scans except during explicitly global phases.

Key Findings
  • Declarative Attention (DA): a prompt-driven, parseable protocol that partitions generation into three modes (global, focus, local). The model emits mode tags in its chain-of-thought and the inference engine updates a segment-level attention mask accordingly, treating tags like tool calls.

  • Zero-shot on off-the-shelf models: DA works without fine-tuning across 15 long-context tasks and on models like Gemma-4-31B and Qwen-3.6-27B.

  • Cost–accuracy trade-off: DA reduced average decoding-attended tokens by ~52.0% for Gemma-4-31B and ~31.1% for Qwen-3.6-27B while incurring small accuracy drops (≈1.27 and 2.75 percentage points) that shrink with model scale.

  • Systems impact: token savings grow with context length (up to millions of tokens per response), and an implementation in vLLM demonstrates practical wall-clock and bandwidth reductions when the middleware can mask KV reads.

Who it's for and trade-offs

Great fit if you operate long-context LLM inference (multi-hundred-thousand+ token contexts) and can modify the serving loop to parse model-declared attention tags; useful for reducing memory bandwidth and decode cost while retaining near-vanilla accuracy. Look elsewhere if you require exact bit-for-bit parity with a baseline (DA incurs modest accuracy changes), cannot alter the inference stack, or expect adversarial contexts where declared scopes may be unreliable. DA is a prompting-level, zero-shot floor — further gains likely with targeted training or tighter model–runtime co-design.

How it works (brief)

A DA state machine watches generated tokens for predefined tags and flips between modes: global (read full context), focus (read one named chunk), and local (read only recent output). The engine then applies a segment-level attention mask to the KV cache (block-aligned, compatible with FlashAttention) so most steps skip large KV reads rather than scanning the entire context.

Information

  • Websitearxiv.org
  • OrganizationsKAIST AI, Google DeepMind
  • AuthorsNamgyu Ho, Huzama Ahmad, Woosung Koh, Se-Young Yun, Tal Schuster, Cicero Nogueira dos Santos
  • Published date2026/09/02

More Items

Generates compact keyword sets for both queries and items with LLMs and matches them directly via an inverted index. Uses supervised fine-tuning to align keyword spaces, then alternates GRPO-based reinforcement learning on query- and item-side generators to co-evolve representations and maximize retrieval F1 while staying compatible with keyword-based infrastructure.

Studies looping shared transformer layers in Mixture-of-Experts models under matched budgets and proposes SMELT: loop the middle half twice while matching per-token FLOPs, non-embedding parameters, and KV cache. Shows 6.8–18.0% training-FLOPs savings on the compute-optimal frontier, stronger downstream gains on code and long-context tasks.

Turns sparse per-student records into individualized simulators that both reproduce a student’s responses and update them under tutor guidance using pooled LLM pretraining followed by per-student specialization; releases StudentSimEval and reference simulators across chess, L2 writing, and math.