On-Demand Attention: Language Models Know When to Recall
Haibo Feng and colleagues show that a pretrained model's decoding states already predict whether a global attention read will help, and use that signal to invoke global attention selectively.
Ask this paper
Only a recall head is trained. Pretrained weights stay frozen and the full historical KV cache remains available, so the method adds a router rather than changing the model.
The predictive signal precedes the read. Local decoding states carry information about the benefit of the global read before it happens, which is what makes skipping it safe.
GPU-side conditional execution in vLLM. The reduced global reads translate into measured decoding speedups over full attention at long context lengths rather than into theoretical savings.
Most of the local-attention loss is recovered. Across Qwen and Gemma models including hybrid-attention backbones, selective recall recovers most of the quality lost under local attention while substantially reducing global reads.
Abstract
Reasoning and agentic workloads increasingly demand efficient long-context inference. Yet full-attention decoding reads the growing history at every step, regardless of its benefit to the next prediction. We show that a pretrained model's decoding states already contain information predictive of this benefit, before the global read. Building on this finding, we introduce On-Demand Attention (ODA), a local-first decoding method that uses a lightweight recall head to selectively invoke global attention as its predicted benefit changes during generation. ODA trains only the recall head, leaving pretrained weights unchanged and the complete historical KV cache available for future recall. We further implement GPU-side conditional execution in vLLM, translating reduced global reads into practical decoding speedups over full attention at long context lengths. Experiments across Qwen and Gemma models, including hybrid-attention backbones, show that selective recall recovers most of the performance lost under local attention while substantially reducing global reads. These findings support long-context inference in which pretrained models guide their own access to the information they retain.