OpenAttention: One Percent Is Enough with Tail Proxies
Sparse attention must approximate the long tail of unseen KV-cache tokens, yet discarding it biases the output and random sampling misses heavy hitters at extreme sparsity. OpenAttention folds the tail into a few count-aware mean proxies discovered via query-geometry clustering, achieving 99% of dense accuracy while reading only 1% of the cache. The accompanying Triton kernel delivers a 10× speedup in the decode attention stage, making ultra-low-density attention practical for long-context inference.
The paper delivers strong empirical results—86.30 mean accuracy on RULER-32K HARD at 1% read versus 87.04 for dense attention—but its theoretical story is more tentative than the abstract suggests. The closed-form analysis relies on unnormalized weight discrepancy $\mathcal{E}$ and Gaussian population models, while the deployed ErrorGate uses a diagonal variance approximation and a proportional budget allocation that the authors explicitly label as heuristic. The contribution is therefore best viewed as a well-engineered system built around a keen empirical observation rather than a method with provable end-to-end guarantees.
The central empirical finding is robust: a single count-aware mean proxy consistently outperforms random sampling at low read densities, both behind an oracle top-$k$ and behind HashAttention. As the authors note, "sampling approaches twice the discard discrepancy when $m \ll N$; this is not a factor-of-two claim for normalized attention or output error" (Section 2.2). The ablations validate the pipeline: QGC improves over key-only $k$-means (86.30 vs. 85.42 in Table 1), QGS matches or exceeds prior selectors in recall at lower latency (Table 2), and the end-to-end result "preserves accuracy close to dense attention at every density" (Section 4.1).
Several important approximations are glossed over in the main narrative. The paper acknowledges that the discrepancy metric $\mathcal{E}$ "diagnoses unnormalised weight distortion but does not by itself explain sampling's density-dependent output failure" (Section 2.2), yet the prominence of the closed-form results risks obscuring this limitation. More critically, the ErrorGate score "is neither an exact finite-cluster discrepancy nor a certified output bound; no independence assumption between weights and values is used" (Section 3.3), and the 10× speedup claim applies only to the decode attention stage at batch 32, not to end-to-end generation. The diagonal variance approximation in ErrorGate "cannot recover the exact error" (Section 6), leaving the method reliant on empirically tuned heuristics for its strongest results.
The evidence strongly supports the claim that proxies surpass sampling for the reported benchmarks, though its breadth is narrow. The RULER-32K HARD experiments focus on seven tasks selected precisely because they stress sparse attention, which amplifies the gap between OpenAttention and sampling-based baselines. The authors verify that vAttention uses sampling "under its authors' per budget protocol" and MagicPIG employs LSH-based sampling with parameters realising the target read (Section 4, Setup). My own check confirms that vAttention indeed combines top-$k$ with random sampling adjusted to an $(\epsilon,\delta)$ tolerance, so the baseline characterisation is fair. Cross-model results on DeepSeek-R1-Distill-Llama-8B and Mistral-7B (Table 5) are encouraging, with "OpenAttention stays within 0.7 points of dense on all three models while the baselines lose up to 29.9 points" (Section 4.3), though still confined to RULER-family tasks; broader real-world workloads remain untested.
Reproducibility is aided by unusually thorough documentation: "The method is defined in §3. QGC is given in Eq. (4), the QGS in Eq. (5), the ErrorGate score in Eq. (6) and the budget allocation in Eq. (7)" (Reproducibility statement), and code is released on GitHub. However, replication is still non-trivial. The 1% configuration relies on multiple interdependent hyperparameters—cluster count $C=64$, water-filling threshold $\lambda$, and a per-budget grid detailed in Appendix D—that are not derived from first principles. Independent verification would require re-implementing the online softmax proxy merge and the int8 QGS coding exactly as in OpenAttention.flash, whose full Triton details are only summarized in Appendix C.
Sparse attention reads only a small subset of the key-value cache, and how to treat the unread tail remains a key challenge. Discarding the tail biases the output, and sampling it as previous approaches do fails at reads of 5% and below, where the tail likely holds heavy hitters the selector missed. Our key observation is that the simplest possible proxy, one token that folds the entire tail into its count-aware mean, already yields lower output error than a sample at every read behind a practical selector. We further show in closed form that at low reads sampling nearly doubles the weight error of discarding while a single proxy stays below it. Building on these findings, we explore approximating the dense attention distribution with a few proxies that target the tail. We propose OpenAttention, a pipeline of clustering, ErrorGate and top-k selection followed by a count-aware mean proxy. To implement clustering and selection both efficiently and effectively, we design training-free Query-Geometry Clustering (QGC) and a Query-Geometry Selector (QGS), which compress the keys while respecting the query directions. Additionally, unlike prior methods that rank clusters by centroid alone, ErrorGate allocates the exact read by the estimated weight error of each cluster. On RULER-32K HARD with Llama-3.1-8B, OpenAttention retains 99% of dense accuracy at a 1% read. It also achieves state-of-the-art accuracy across models, context lengths and datasets. Our customized attention kernel, OpenAttention.flash, runs the decode attention stage 10$\times$ faster than dense attention at a 1% read.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.