An Everyday Analogy
Imagine writing a long email. Most of the time your attention stays on the current paragraph—the previous sentence just finished, the next one flows naturally. But occasionally you need to scroll back to check a detail: what was that number? How is that name spelled?
Re-reading the whole email after every word would be far too slow. But never looking back means you may contradict yourself by the end.
This is exactly the dilemma facing long-context LLMs.
- Full Attention: re-attends over the entire context for every generated token. Accurate, but slow—cost per step is O(n) in context length.
- Local Attention: only sees the most recent 2048 tokens. Fast, but prone to "amnesia"—important earlier information may be missed.
- On-Demand Attention (ODA): do the local computation first, then let a lightweight "recall head" decide: "does this step need to look back?" If yes, run full attention; if no, use the local result.
- Full: 37.94
- ODA: 36.82 (92.2% of the gap recovered, 70.6% Full call rate)
- Local: far lower
- Qwen3-8B: 91.07 (Full 92.59, Local 25.43), 47.3% Full call rate
- Qwen3.5-2B (hybrid architecture): 94.08 (Full 94.35), 44.3% call rate—only 6 full-attention layers are gated; the other 18 Gated DeltaNet layers keep native computation
- Gemma-4-12B-it: also effective across model families
- ODA's FLOPs are only 24.26% of Full—a 75.74% saving, a 4.12x compute ratio
- Measured in vLLM, throughput rises from 75.54 to 149.52 tokens/s—a 1.98x speedup
- Among 675 positions where Local predicted incorrectly, the recall head's benefit-sign AUROC is 0.643
- Local's entropy scores only 0.486 AUROC
- Using only H+E (history + current token): Full is called almost always (99.73%)—the head never learns when *not* to call
- Adding L (the current Local result): the call rate drops to 52.58%—the Local computation's output is the key signal
- No modification of main model weights—the pretrained model is fully frozen
- No architecture change—just an external small head
- Full KV cache preserved—Full calls access the complete history, not a compressed or truncated version
- Learns "timing", not "scope"—deciding *when* to look, not *how much* to look at
- Not suited to short contexts: below 4K, ODA is slower than Full—the local-branch overhead isn't justified
- The recall head requires training: small (28.3M), but needs 196K paired training samples
- Only base models tested: post-training behavior was not evaluated
- Threshold tuning needed: default θ=0, but different applications may require different thresholds
- Judge: the recall head predicts "would full attention help?"
- Gate: decide whether to execute the expensive full-attention computation
- Paper: https://arxiv.org/abs/2609.20734
The Core Question: Can a Model Predict What It Needs?
The idea sounds simple, but there's a critical premise: before deciding whether to invoke full attention, the model must already "know" whether full attention would help.
How is that possible before the full-attention computation has run?
The paper's insight: the model's decoding state itself contains this information.
Concretely, the recall head reads three signals:
1. The current token's embedding (E(x_t)) — what is being processed now 2. The current Local computation's hidden state (h_t^L) — what the local view sees 3. The previous step's hidden state (h_{t-1}) — accumulated prior context
These three signals pass through a small 28.3M-parameter head (only 1.7% of Qwen3-1.7B's 1.7B parameters) and output a scalar score q_t. If q_t exceeds a threshold θ (default 0), full attention is triggered.
Training: Paired Supervision
How is the training data built? Via paired computation.
At each position during training, both Local and Full computations are run, and their predicted probabilities of the correct next token are compared. If Full is significantly better than Local, the recall head learns to say "yes"; if similar, "no".
The objective uses Huber regression (robust to large errors) plus a cost penalty λ—each Full call has a cost, so it's only worth invoking when the benefit exceeds the price.
Training data: 196,608 samples, trained on Qwen3-1.7B. The main model weights are completely frozen; only the recall head is trained.
Key Results
Quality Recovery
On Qwen3-1.7B:
| Strategy | RULER16K score | Full call rate | |---|---|---| | Full | 81.94 | 100% | | ODA | 81.17 | 41.6% | | Local | 19.23 | 0% |
ODA recovers 98.8% of the gap between Full and Local, using only 41.6% full-attention calls.
On LongBench v1:
Cross-Model Generalization
The same approach works across models:
Learned Timing vs. Random Calls
This is the most crucial ablation: what matters isn't how often Full is called, but when.
Across 5 RULER16K tasks, ODA scores 88.96 at a 41.33% call rate; random calls at a matched 40.99% rate score only 32.79±0.95. The Full baseline is 88.63.
In other words: calling Full 41% of the time with learned timing nearly matches 100% Full, while random calls perform about as poorly as Local.
This proves the recall head learns meaningful "timing judgment", not random guessing.
Compute Savings
At 128K context and a 12.5% Full call rate:
Caveat: at short 4K contexts, ODA is actually 15.8% *slower* than Full—the local-branch overhead isn't worth it. ODA is a long-context-specific optimization.
Deeper Insights
1. "Correctness" and "Access Benefit" Are Different Things
A subtle finding: among 1,342 positions where Local already predicted the top-1 token correctly, Full raised that token's probability at 718 positions but *lowered* it at 437.
That is: even when Local guesses the right token, Full may still offer a better probability distribution. Conversely, even when Full yields positive NLL gain, it may not change the top-1 prediction.
Recall isn't just "fixing errors"—it's optimizing the whole distribution, not just top-1 accuracy.
2. The Recall Head Isn't Learning "Uncertainty"
A natural hypothesis: the head merely detects where "Local is uncertain". But the paper's controls show otherwise:
The learned signal differs from simple uncertainty measures—it estimates "how much full attention would improve things", not "how uncertain Local is".
3. Input Ablations Reveal the Decision Mechanism
This means: the crucial information for "do I need to look back?" comes precisely from "taking a quick glance first". You must run the Local computation before you can know whether Full is needed.
This matches human reading: you skim the current sentence first; only if it seems to contradict something earlier do you flip back. Not scrutinizing every sentence, never ignoring them entirely—glance first, then decide.
4. How ODA Differs from Other Conditional-Compute Methods
Compared with prior approaches (CoLT5, AHA, L2A):
Limitations
Implications
Another Case of "Judge-Gate Decoupling"
ODA is a textbook instance of the judge–gate pattern:
This parallels earlier examples like DoM (hidden-state detection of reward hacking) and NavTrust (benchmark scores predicting real reliability): a lightweight, specialized judge decides whether to trigger expensive computation.
"Glance First, Decide Later" Is a General Pattern
ODA's "Local first, then decide on Full" strategy is isomorphic to speculative decoding (small model drafts, large model verifies). This "cheap probe first, deepen on demand" pattern is a nearly universal optimization for compute-constrained settings.
"The Model Knows What It Needs"
The deepest insight: the model's intermediate states already encode "what I need". The recall head simply extracts it. This aligns with "linear probe" findings in interpretability—hidden states contain richer, more reliable information than CoT reveals.
One-sentence summary: On-Demand Attention trains a 28.3M-parameter recall head that recovers 98.8% of long-context performance on Qwen3-1.7B with only 41.6% full-attention calls, achieving a 1.98x decoding speedup at 128K context. The core idea: run local computation first, then use its result to decide whether full attention is needed—mirroring the human "glance first, look back when necessary" reading strategy.