Openattention achieves high accuracy reading just one percent of data
OpenAttention: One Percent Is Enough with Tail Proxies
Machine LearningArtificial Intelligence
Summary
Attention systems in AI models usually look at a lot of data, but that is slow and expensive. This paper shows that reading just one percent of data with a smart approach called OpenAttention can still give nearly the same results. The key idea is to summarize unread parts with a single proxy token rather than sampling small samples that might miss important info. The authors created new ways to cluster and select data efficiently without additional training. Their method runs much faster while keeping output quality high.
What this means in practice
- •For natural language processing teams: Speed up transformer-based language model decoding with minimal accuracy loss using OpenAttention's efficient 1% read mechanism.
- •For machine learning infrastructure engineers: Implement faster attention kernels to reduce computational costs in large-scale AI model serving without retraining.
Authors
Hyunwoo Yu, Jongbeom Lee, Hyunbin Kim, Sungjun Park, Jaemin Choi, Jincheol Yang, Seunghun Moon, Bongjoon Hyun, Kyomin Sohn, Kyeongbo Kong, Suk-Ju Kang
Abstract
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.