arXiv preprint ·

AttentionIndex: Accelerating Learned Sparse Attention with a Causal Inverted-File Index

Jeongmin Brian Park and Wen-mei Hwu

Published
Source
arXiv preprint, posted before peer review
Original
Read it at arXiv
Topics
Artificial IntelligenceMachine Learning

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.

The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.

Read it on ReadSearch

Each opens after you create an account, and brings you back to this paper.

Already a reader? Sign in

All new papers in Artificial Intelligence

Get the few papers worth your time in Artificial Intelligence every morning.

ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.

Papers belong to their authors and publishers; every page here links to the original. All research topics