On-Demand Attention: Teaching Models When They Need to Look Back

A Daily Analogy

Imagine writing a long email. Most of the time your attention stays on the current paragraph—the previous sentence is done, the next one follows naturally. But occasionally you need to scroll back to check a detail: what was that number? How is that name spelled?

Re-reading the entire email after every word would be far too slow. But never looking back means you may contradict yourself later.

This is exactly the dilemma facing long-context LLMs:

• Full Attention: every generated token re-attends to the entire context. 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 "forgetting"—important earlier information can 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.

The Core Question: Can a Model Predict What It Needs?

The idea sounds simple but has a crucial precondition: the model must already "know" whether full attention would help—before running it.

How is that possible? 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 this step is processing 2. The hidden state of the current Local computation (h_t^L) — what the local view sees 3. The previous step's hidden state (h_{t-1}) — accumulated prior context

These three signals feed a small 28.3M-parameter head (1.7% of Qwen3-1.7B's size) that outputs a scalar score q_t. If q_t exceeds a threshold θ (default 0), full attention is triggered.

Training: Paired Supervision

How is training data built? Via paired computation: at each position, run both Local and Full, then compare their predicted probabilities of the correct next token. If Full is significantly better than Local, the recall head should learn 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 invoked when the benefit exceeds the cost.

Training data: 196,608 samples, trained on Qwen3-1.7B. The base 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 while making only 41.6% of full-attention calls.

On LongBench v1:

• Full: 37.94
• ODA: 36.82 (92.2% of the gap recovered, 70.6% call rate)
• Local: far lower

Cross-Model Generalization

The same approach works across models:

• Qwen3-8B: 91.07 (Full 92.59, Local 25.43), 47.3% 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: works across model families too

Learned Timing vs. Random Calls

This is the most critical control experiment: it's not how often Full is called that matters—it's when.

Across 5 RULER16K tasks, ODA scores 88.96 with a 41.33% call rate; random calls at an identical 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 calling is barely better than Local. The recall head has genuinely learned meaningful timing, not random guessing.

Compute Savings

At 128K context with a 12.5% Full call rate:

• ODA uses only 24.26% of Full's FLOPs—a 75.74% saving, a 4.12x compute ratio
• In vLLM measurements, throughput rises from 75.54 to 149.52 tokens/s—a 1.98x speedup

Caveat: at short 4K contexts, ODA is actually 15.8% slower than Full—Local's overhead isn't worth it. ODA is a long-context optimization.

Deeper Insights

1. "Correctness" and "Access Benefit" Are Different Things

A subtle finding: among 1,342 positions where Local already predicted the correct top-1 token, Full *increased* that token's probability at 718 positions and *decreased* it at 437.

That is: even when Local guesses the right token, Full can still provide a better probability distribution. Conversely, a positive NLL benefit from Full doesn't necessarily change the top-1 prediction. Recall isn't just "fixing errors"—it optimizes the whole distribution, not just top-1 accuracy.

2. The Recall Head Doesn't Learn "Uncertainty"

A natural hypothesis: the head merely detects positions where Local is uncertain. But controls show otherwise:

• Among 675 positions where Local predicted incorrectly, the head's benefit-sign AUROC is 0.643
• Local's entropy AUROC is only 0.486

The signal the head learns is different from simple uncertainty measures—it learns "how much improvement full attention would bring," not "how uncertain Local is."

3. Input Ablations Reveal the Decision Mechanism

• Using only H+E (history + current token): almost always calls Full (99.73%)—it never learns "when not to"
• Adding L (the current Local result): the call rate drops to 52.58%—the Local computation is the key signal

This means: the critical information for deciding whether to look back comes precisely from glancing first. You must run Local before knowing whether Full is needed. This matches human reading: skim the current sentence; if it seems to conflict with earlier text, then flip back. Not careful reading of everything, not ignoring everything—look fast first, then decide.

4. How ODA Differs from Other Conditional Computation

Compared to prior methods (CoLT5, AHA, L2A):

• Doesn't modify base model weights—the pretrained model is fully frozen
• Doesn't modify the architecture—only an external small head is added
• Keeps the full KV cache—Full calls access the complete history, not a compressed or truncated version
• Learns "when" not "what range"—it decides when to look, not how much

Limitations

• Not suited to short contexts: below 4K, ODA is slower; Local overhead isn't worth it
• The recall head requires training: small (28.3M), but needs 196K paired samples
• Only base models tested: post-training behavior not evaluated
• Threshold tuning needed: default θ=0, but different applications may need different values

Takeaways

Another Instance of "Judge-Gate Decoupling"

ODA's architecture exemplifies judge-gate decoupling: a judge (the recall head predicting whether full attention helps) controls a gate (whether the expensive computation executes). This parallels using hidden states to detect reward hacking, or benchmark scores to predict real-world reliability: a lightweight, specialized judge decides whether to trigger expensive computation.

"Look Fast First, Then Decide" Is a General Pattern

ODA's "Local first, then decide on Full" mirrors speculative decoding (small model drafts, large model verifies). "Try cheaply first, then go deep on demand" is a near-universal optimization pattern for compute-constrained settings.

"The Model Knows What It Needs"

The deeper insight: the model's intermediate states already contain the information about what it needs. The recall head merely extracts it. This aligns with interpretability findings from linear probes—hidden states hold more reliable information than CoT reveals.

• Paper: https://arxiv.org/abs/2609.20734

One-sentence summary: On-Demand Attention trains a 28.3M-parameter recall head that, on Qwen3-1.7B, recovers 98.8% of long-context performance with only a 41.6% full-attention call rate and achieves a 1.98x decoding speedup at 128K context. The core insight: compute locally first, then use the local result to decide whether full attention is needed—mirroring the human reading strategy of "glance first, look back only when necessary."