Language models learn to recall past information only when needed
On-Demand Attention: Language Models Know When to Recall
Computation and Language
Summary
Long language conversations or documents are hard for AI to read efficiently because they look at everything every time. The authors found that language models already have clues inside them that tell when looking at the whole past will help predict the next word. They built a system called On-Demand Attention that uses a small part of the model to decide if it should check the entire history or just a nearby part. This way, the system works faster without losing much accuracy, making AI better at reasoning over long texts.
What this means in practice
- •For language model engineers: Speed up inference of large language models on very long inputs by reducing unnecessary attention computations.
- •For natural language processing teams: Improve efficiency in tasks requiring reasoning over long documents by selectively recalling relevant context.
Authors
Haibo Feng, Ruiqi Liang, Hanyang Peng, Shiqi Yu
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.