On-Demand Attention: Language Models Know When to Recall
Introduces On-Demand Attention (ODA) method, significantly reducing global reads and enhancing long-context inference efficiency.
Key Findings
Methodology
This study introduces a novel On-Demand Attention (ODA) method, using a lightweight recall head to selectively invoke global attention during generation based on predicted benefits. ODA trains only the recall head, keeping pretrained weights unchanged, with the complete historical KV cache available for future recall.
Key Results
- On the RULER16K dataset, ODA uses global attention on 41.6% of steps, improving the score from 19.23 under local attention to 81.17, close to the full attention score of 81.94.
- On the Qwen3-1.7B model, ODA reduces major decoding operation FLOPs by approximately 76% and achieves 1.98× the single-request decoding throughput.
- On the Qwen3.5-2B hybrid architecture, ODA limits access in its six full-attention layers, scoring 94.08 compared to 94.35 with full attention.
Significance
This study significantly enhances the efficiency of long-context inference by reducing global reads, supporting pretrained models in self-guiding information access in long contexts. This method offers faster decoding speeds and lower computational costs without altering pretrained weights, making it highly significant for long-context inference tasks in academia and industry.
Technical Contribution
The ODA method achieves on-demand invocation of global attention through a lightweight recall head without altering pretrained model weights, offering new engineering possibilities and theoretical guarantees. Unlike existing dynamic attention methods, ODA leverages existing decoding states to predict the relative benefit of global attention.
Novelty
ODA is the first to achieve on-demand invocation of global attention based on predicted benefits using existing decoding states without altering pretrained weights.
Limitations
- In some tasks, selective recall may lead to slight performance drops, especially when frequent global context is needed.
- Training the recall head requires additional computational resources.
Future Work
Future work could explore applications on larger-scale models and more tasks, as well as further optimization of recall head training efficiency.
AI Executive Summary
As the demand for long-context inference in natural language processing tasks increases, existing full-attention decoding methods are inefficient in handling long contexts. Full-attention decoding reads the growing history at every step, regardless of its actual benefit to the next prediction.
This paper introduces a novel On-Demand Attention (ODA) method, using a lightweight recall head to selectively invoke global attention during generation based on predicted benefits. This method trains only the recall head, keeping pretrained weights unchanged, with the complete historical KV cache available for future recall. Experiments show that ODA significantly reduces global reads while recovering most of the performance lost under local attention in Qwen and Gemma models.
This study provides a new solution for long-context inference, supporting pretrained models in self-guiding information access in long contexts, significantly enhancing decoding efficiency and speed. This method offers faster decoding speeds and lower computational costs without altering pretrained weights, making it highly significant for long-context inference tasks in academia and industry.
Deep Analysis
Background
As the demand for long-context inference in natural language processing tasks increases, existing full-attention decoding methods are inefficient in handling long contexts. Full-attention decoding reads the growing history at every step, regardless of its actual benefit to the next prediction.
Core Problem
Existing full-attention decoding methods are inefficient in long-context inference because they read the growing history at every step, regardless of its actual benefit to the next prediction.
Innovation
This paper introduces a novel On-Demand Attention (ODA) method, using a lightweight recall head to selectively invoke global attention during generation based on predicted benefits. This method trains only the recall head, keeping pretrained weights unchanged, with the complete historical KV cache available for future recall.
Methodology
- �� Use a lightweight recall head to predict the relative benefit of global attention.
- �� Selectively invoke global attention based on predicted benefits.
- �� Keep pretrained weights unchanged, with the complete historical KV cache available for future recall.
Experiments
Experiments were conducted on Qwen and Gemma models using RULER16K and LongBench datasets to evaluate the effectiveness of ODA in reducing global reads and recovering performance.
Results
On the RULER16K dataset, ODA uses global attention on 41.6% of steps, improving the score from 19.23 under local attention to 81.17, close to the full attention score of 81.94.
Applications
This method can be applied to natural language processing tasks requiring long-context inference, such as machine translation and text generation.
Limitations & Outlook
In some tasks, selective recall may lead to slight performance drops, especially when frequent global context is needed.
Plain Language Accessible to non-experts
Imagine a librarian managing books. Every time someone borrows a book, the librarian checks all book records, which is inefficient. On-Demand Attention is like the librarian only checking relevant book records each time a book is borrowed, thus improving efficiency. The librarian uses a lightweight tool to predict which book records are relevant and then only checks those records. This is like the librarian selectively checking book records based on need during the borrowing process, instead of checking all records every time.
ELI14 Explained like you're 14
Imagine you work in a school library, and every time someone borrows a book, you have to check all the book records, which is inefficient. On-Demand Attention is like only checking the relevant book records each time a book is borrowed, thus improving efficiency. You use a small tool to predict which book records are relevant and then only check those records. This is like selectively checking book records based on need during the borrowing process, instead of checking all records every time.
Glossary
On-Demand Attention
A method that selectively invokes global attention during generation based on predicted benefits.
Used to enhance long-context inference efficiency.
Recall Head
A lightweight component used to predict the relative benefit of global attention.
Used in the ODA method to selectively invoke global attention.
Global Attention
An attention mechanism that reads the growing history at every step.
Used for long-context inference.
Local Attention
An attention mechanism that reads history only within a fixed window.
Used to reduce computational costs.
KV Cache
A cache that stores historical keys and values.
Used to reuse historical information during decoding.
Open Questions Unanswered questions from this research
- 1 How to apply the ODA method to larger-scale models to further enhance long-context inference efficiency.
- 2 How to optimize the training efficiency of the recall head to reduce computational resource consumption.
Applications
Immediate Applications
Machine Translation
In machine translation tasks, using the ODA method can improve translation efficiency and accuracy.
Long-term Vision
Smart Assistants
Applying the ODA method in smart assistants can enhance understanding and response capabilities in long conversations.
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.