🤖 machine learning

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers

SpotAttention is a lightweight, plug-in block-sparse routing mechanism that attaches to frozen pretrained transformers to learn attention distributions via KL distillation, enabling accurate long-context inference up to 128K tokens with significantly faster decoding speeds and reduced memory usage compared to existing baselines.

Original authors: Huzama Ahmad, Se-Young Yun

Published 2026-06-23

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Huzama Ahmad, Se-Young Yun

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a very smart librarian (the AI model) who has read millions of books. When you ask a question, the librarian usually tries to remember every single page they've ever read to find the answer.

The problem is that as the library grows, this "remembering everything" approach becomes incredibly slow and expensive. It's like trying to find a specific needle in a haystack that keeps getting bigger every second.

SpotAttention is a new, lightweight "assistant" that helps the librarian find the right pages without having to read the whole library again. Here is how it works, broken down into simple concepts:

1. The Problem: The "Heavy Backpack"

When an AI tries to answer a question based on a very long document (like a whole novel or a massive codebase), it has to keep a "backpack" of all the important words it has seen so far.

The Old Way: Every time the AI thinks of a new word, it dumps the entire backpack onto the table to look for clues. As the document gets longer, the backpack gets heavier, and the table gets too crowded. This makes the AI slow and expensive to run.
The "Sparse" Idea: Other researchers tried to say, "Hey, the AI only really needs to look at a few specific pages, right? Let's just look at those." But figuring out which pages to look at was itself a slow, expensive task. It was like hiring a new person just to decide which pages to read, and that person was almost as slow as reading the whole book.

2. The Solution: The "Spotter" (SpotAttention)

The authors created SpotAttention, which is like a tiny, super-fast "Spotter" that attaches to the librarian.

How it learns: The Spotter doesn't need to be taught from scratch. It watches the "expert" librarian (the pre-trained AI) work. It learns to mimic the expert's eye, figuring out which pages the expert would naturally look at.
The Magic Trick: Instead of reading every single word to decide what to keep, the Spotter looks at chunks of text (like paragraphs) and quickly scores them. It's like a librarian who can glance at a bookshelf and instantly know, "I only need to pull out these three books; the rest can stay on the shelf."

3. The "Dual Top-p" Rule: A Smart Budget

The paper introduces a clever rule called Dual Top-p. Imagine the librarian has a budget for how many pages they can look at.

The Old Way: Some systems say, "You can only look at the top 50% of pages." This is rigid. Sometimes you need 80%, sometimes 20%.
The SpotAttention Way: The Spotter looks at the "importance" of the pages and says, "Okay, for this specific question, we need to keep the first few pages (the start), the last few pages (what was just said), and then the most important middle pages."
It dynamically adjusts the budget. If the question is simple, it grabs fewer pages. If it's complex, it grabs more. It does this automatically, without needing a second, slow step to decide.

4. The Results: Speed and Memory

The paper tested this on several large AI models (specifically the Qwen family) with very long contexts (up to 128,000 words).

Speed: At a length of 128,000 words, SpotAttention was 3.9 times faster than the standard method (FlashAttention) and 1.8 times faster than the current best alternative (Twilight).
Accuracy: Despite skipping most of the text, the AI got the answers just as right as if it had read everything. It didn't lose any "intelligence."
Memory: The Spotter's "notebook" (the cache it uses to make decisions) can be shrunk down to a tiny size (using special compression) without losing accuracy. This saves a lot of computer memory.

5. Why It's Different

No Re-training: You don't have to re-teach the whole AI. You just plug this tiny Spotter onto an AI that is already finished and working.
It's Learned, Not Hard-Coded: Previous methods used fixed rules (like "always skip the middle"). SpotAttention learns the pattern of what is important, making it smarter and more flexible.
It Works Everywhere: They tested it on different sizes of AI models (from small 4-billion parameter models to large 32-billion ones) and it worked well on all of them.

In summary: SpotAttention is a lightweight, learned assistant that helps AI models read long documents by quickly identifying the most important parts to focus on. It makes the AI significantly faster and cheaper to run on long texts without making it any less smart.

Technical Summary: SpotAttention

Problem Statement

Long-context capabilities have become standard in pretrained Large Language Models (LLMs), enabling tasks such as processing entire codebases or grounding answers in extensive evidence. However, the computational cost of self-attention scales poorly with sequence length. During the prefill phase, compute grows quadratically with sequence length, and during the decode phase, every new token requires re-reading a Key-Value (KV) cache that grows linearly with context size. While FlashAttention improved I/O efficiency by avoiding materializing the full attention matrix, it did not alter the fundamental quadratic prefill cost or the linear growth of per-step memory reads.

Existing sparse attention methods attempt to mitigate this by attending only to a relevant subset of past tokens. However, selecting this subset is often expensive. Methods like Quest score token blocks, and Twilight adds top- $p$ pruning, but the selector itself becomes a new bottleneck. Furthermore, training-free selectors often rely on hand-designed heuristics that may not generalize well or require multiple stages of selection, incurring cumulative overhead.

Methodology: SpotAttention

SpotAttention is a plug-in, block-sparse routing mechanism designed to retrofit onto already-trained, frozen transformers. It consists of a lightweight "selector" (or indexer) attached to every full-attention layer.

1. Selector Architecture

The selector is a tiny multi-head Q-K scorer with reduced width ( $d_{idx}$ ) and $H_{idx}$ heads. Instead of scoring individual tokens, it operates on blocks of queries and keys (e.g., 16 tokens per block).

Scoring Mechanism: For a query at position $q$ and a key block at position $t$ , the selector computes a head-weighted ReLU score:
$s_s[q, t] = \sum_{h=1}^{H_{idx}} w_h \text{ReLU}(\mathbf{q}^I_h \cdot \mathbf{k}^I_{h,t})$
where $\mathbf{q}^I_h$ and $\mathbf{k}^I_{h,t}$ are per-head projections, and $w_h$ are mixing weights derived from the hidden state.
Efficiency: By scoring block-wise, the selection process runs as a single tensor-core matrix multiplication, avoiding the overhead of token-level operations.

2. Training via KL Distillation

The selector is trained independently against the dense attention distribution of the frozen backbone (the "teacher").

Objective: The training minimizes the Kullback-Leibler (KL) divergence between the teacher's attention distribution ( $p_\tau$ ) and the selector's estimated distribution ( $p_\sigma$ ).
Variants: The paper explores two KL variants:
- DenseKL: Computes loss over all $T$ keys.
- SparseKL: Computes loss only over the selector's own top- $K$ selected set, with distributions renormalized.
Convergence: The selector converges rapidly (within ~200 steps) across various backbone sizes (4B to 32B) and architectures, learning to match the teacher's distribution shape.

3. Dual Top- $p$ Selection Rule

At inference, the selector scores all past blocks. To determine which blocks to attend to, SpotAttention introduces a dual top- $p$ rule that reads the budget directly from the selector's calibrated distribution:

Reserved Blocks: A fixed prefix of "sink" blocks (e.g., first 128 tokens) and a fixed suffix of "recency" blocks (e.g., last 256 tokens) are always retained to prevent them from being dominated by the nucleus selection.
Nucleus Selection: Top- $p$ (nucleus) sampling is applied to the residual distribution (blocks excluding sinks and recency). The selection stops when the cumulative probability mass reaches a threshold $p$ .
Dynamic Budget: This approach yields a variable budget per query and per layer, adapting to the content without a separate pruning stage. A minimum floor ( $K_{min}$ ) prevents degenerate cases where the nucleus is too small.

4. Quantization

The selector's K-cache can be quantized to INT4 or FP4 microscale. Since the selector's cache is distinct from the backbone's KV cache, this reduces the memory footprint of the selector by approximately 3.5 $\times$ without affecting the backbone's precision or accuracy.

Key Contributions

Plug-In Block-Sparse Attention: A lightweight selector trained via KL distillation that attaches to frozen transformers. It runs as a single tensor-core matmul, avoiding the heavy selector overhead seen in prior work.
Dual Top- $p$ Rule: A content-adaptive selection mechanism that derives the per-query, per-layer budget directly from the learned distribution. It reserves sink and recency blocks to ensure stability, matching static top- $K$ accuracy with a smaller mean budget.
Training and Inference Behavior: The paper demonstrates that the selector learns the per-layer budget asymmetry of the backbone (early layers attend to more context than later layers), a pattern static $K$ cannot replicate. The method generalizes from a 16K training length to 128K evaluation without accuracy degradation.

Results

The method was evaluated on five backbones (Qwen3: 4B–32B; Qwen3.5: 4B–9B) across four benchmarks (RULER, BABILong, InfiniteBench, LongBench-v2) at context lengths up to 128K.

Accuracy: SpotAttention matches the dense backbone's accuracy at context lengths up to 128K (8 $\times$ the training length). At a budget of $K=0.5L$ , it stays within one bootstrap standard error of the dense baseline across all models.
Latency and Throughput:
- At $L=128K$ on Qwen3-8B, SpotAttention runs 3.9 $\times$ faster than FlashAttention and 1.8 $\times$ faster than Twilight (the strongest training-free baseline).
- Decode throughput is sustained at 70 tokens/s at 128K and 36 tokens/s at 512K, whereas FlashAttention and Twilight run out of memory (OOM) at 512K and 256K respectively.
- The selector itself is significantly faster than Twilight's selector (1.55 ms vs 4.59 ms at 128K).
Quantization: Quantizing the selector's K-cache to INT4 or FP4 shrinks its memory footprint by 3.5 $\times$ with no measurable accuracy loss. While this does not improve single-stream decode latency (dominated by backbone KV bandwidth), it significantly reduces memory usage and prefill costs.

Significance and Claims

The paper positions SpotAttention as a solution to the "selector bottleneck" in sparse attention. By learning a calibrated distribution via KL distillation, the selector enables a dynamic, nucleus-based budget that adapts to the specific needs of each query and layer.

The authors claim that:

Efficiency: It achieves the highest throughput among compared methods at long contexts (up to 512K), overcoming memory limitations that halt dense and other sparse baselines.
Generalization: A selector trained on a fixed 16K context generalizes effectively to 128K and beyond without retraining.
Modularity: It is a "plug-in" solution that does not require retraining the backbone, making it applicable to existing pretrained models.
Limitations: The authors note that the largest evaluated backbone is 32B, and results are currently limited to the Qwen family on English benchmarks. They also acknowledge that sparse-prefill latency was not benchmarked, though the architecture supports it, and that the dual top- $p$ rule relies on the selector being well-calibrated (which holds under SparseKL training but collapses under DenseKL).