On-Demand Attention

Language Models Know When to Recall

arXiv:2609.20734 | English | 简体中文

On-Demand Attention (ODA) runs Local attention first, then uses a lightweight recall head to decide whether to recompute the step with Full attention. Only the recall head is trained; the language model stays frozen, and the complete historical KV cache remains available for future recall.

Figure 1. ODA decoding and execution efficiency. Efficiency panels use fixed continuations and prescribed Full-call rates.

Table 1. RULER results from the paper. Percentages below ODA scores are Full-call rates; Avg averages scores and rates separately over the five lengths.

Experiment settings and reproduction status.

This release includes Router training, Transformers inference, evaluation, and experimental vLLM backends for Qwen3/Qwen3.5. Gemma uses Transformers; see the supported models and runtime limits. Inference/evaluation batch size and training microbatch per GPU default to 1.

Trained Router weights will be released separately on Hugging Face. Until then, use an export from your own training run.

git clone https://github.com/Fenghaibo-git/On-Demand-Attention.git
cd On-Demand-Attention
export PYTHONPATH="$PWD/src:$PWD/scripts${PYTHONPATH:+:$PYTHONPATH}"

Install the required environment, then run Transformers inference with a local base model and a trained Router export:

python scripts/generate.py --backend hf \
  --model /path/to/models/Qwen3.5-2B \
  --router /path/to/trainable_state.pt \
  --prompt "Explain how attention works." --max-new-tokens 256

Training · vLLM inference · Evaluation · RULER16K workflow

GPL-3.0. Upstream copyright and license notices are retained; see NOTICE.