On-Demand Attention: Language Models Know When to Recall
On-Demand Attention uses a lightweight recall head to selectively invoke global attention, cutting long-context decoding costs in vLLM with minimal quality loss.
Researchers introduce On-Demand Attention (ODA), a local-first decoding method in which a lightweight recall head predicts when global attention benefits the next token. Only the recall head is trained, leaving pretrained weights and the complete KV cache unchanged. GPU-side conditional execution implemented in vLLM converts reduced global reads into practical decoding speedups at long context lengths. Experiments across Qwen and Gemma models, including hybrid-attention backbones, show selective recall recovers most of the performance lost under local attention.
- Decoding states predict whether global attention helps the next prediction
- Trains only the recall head; pretrained weights stay frozen
- Conditional execution implemented in vLLM for practical decoding speedups
- Validated on Qwen and Gemma models including hybrid-attention backbones
Full article165 words · extracted from arxiv.org · click to collapse
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.
Text extracted automatically; images, tables and formatting may be missing. Original: https://arxiv.org/abs/2609.20734