AttentionIndex boosts long text AI speed and accuracy
AttentionIndex: Accelerating Learned Sparse Attention with a Causal Inverted-File Index
Artificial IntelligenceMachine Learning
Summary
Working with very long texts in AI is slow because the system has to check too many pieces of information. The authors created AttentionIndex, a smarter way to quickly find the most important information by grouping similar parts and focusing only on those. This method speeds up the process while improving the accuracy of understanding long texts. It helps AI models handle millions of tokens much faster without losing quality.
What this means in practice
- •For machine learning engineers: Increase efficiency and accuracy when processing very long text sequences in AI models by using a new indexing method for sparse attention.
- •For search engine developers: Speed up large-scale document retrieval by routing queries through similarity-based inverted indexes to reduce computation without sacrificing relevance.
Authors
Jeongmin Brian Park, Wen-mei Hwu
Abstract
Learned sparse attention, used in production models such as DeepSeek-V4 and GLM-5, lets each query attend to only k keys chosen by a lightweight indexer, but the indexer still scores every key in the prefix. At one million tokens, this selection takes more than half of prefill time. Two-stage selection over contiguous blocks, as in HiSA, reduces this cost, but its recall declines as context grows at a fixed candidate budget. We present AttentionIndex, which routes each query over an inverted-file index that groups indexer keys by similarity and refines the candidates with the native indexer score. AttentionIndex builds the index during prefill, making new keys searchable immediately, reusing previous clustering to keep maintenance cheap, and exposing each query only to clusters built from its own prefix. On GLM-5.3 and DeepSeek-V4-Flash at one million tokens, AttentionIndex delivers 1.3-1.8 times the prefill throughput of exhaustive selection, and at throughput within 5% of HiSA, it recovers 17.0 and 10.0 percentage points more of the exhaustive top-k, respectively. AttentionIndex also improves accuracy on long-context benchmarks by up to 18.4 percentage points over HiSA at the same candidate budget.