提出近似紧致的流式注意力空间下界,解决长期存在的理论缺口。
Towards Tight Bounds for Streaming Attention
- 融合三种核密度估计方法构造紧凑缓存
- 首次实现空间复杂度近乎理论最优
- 适用于低内存场景下的Transformer加速
注意力机制是现代Transformer架构的核心,但其表达能力带来二次时间复杂度和线性空间占用。经典Transformer需存储所有历史输入(标记)以生成下一个输出。在有限内存下实现Transformer的问题,即KV缓存压缩,近年备受关注,催生出诸多启发式方法。Haris等(COLT'25)与Kochetkova等(NeurIPS'25)将此问题形式化为流式注意力近似,并基于分歧理论给出上界、信息论给出下界,但两者间仍存在显著差距。例如,其算法空间随精度参数增大,而下界未相应增强。本文重新审视该问题,提供近乎紧致的时空复杂度下界。算法侧通过三种核密度估计方法——分歧型核样本构造(Charikar-Kapralov-Waingarten'24)、多项式法(Greengard-Rokhlin'87, Alman-Song'23)与空间划分(Andoni-Laarhoven-Razenshteyn-Waingarten'17, Charikar-Kapralov-Nouri-Siminelakis'20)——的巧妙结合实现。下界侧,核心技术贡献为一种利用大量辅助信息的索引问题新方法,或可推广至其他高维几何估计问题。
原文摘要 · Abstract (English)
The attention mechanism is a cornerstone of modern transformer architectures. However, its expressive power comes at the cost of quadratic runtime and linear space usage. In particular, the classical transformer architecture explicitly stores all previously seen input elements (tokens) in order to generate the next one. The problem of implementing a transformer in limited space, known as KV cache compression, has received much interest over the past few years, spurring the development of powerful heuristics. Recent works of Haris et al, COLT'25 and Kochetkova et al, NeurIPS'25, formalized KV cache compression as the streaming attention approximation problem, providing both upper bounds (based on discrepancy theory) and information theoretic lower bounds. However, those papers left open a significant gap between the upper and lower bounds. For example, the space usage of their algorithms increases with the precision parameter, but the lower bound does not get stronger. In this work, we revisit the streaming attention approximation problem and provide nearly tight bounds on its space complexity. On the algorithmic side, we achieve the result through a surprisingly tight interplay between three distinct methods for kernel density estimation: discrepancy-based coreset constructions (e.g., Charikar-Kapralov-Waingarten'24), the polynomial method (e.g., Greengard-Rokhlin'87, Alman-Song'23), and space partitioning (e.g., Andoni-Laarhoven-Razenshteyn-Waingarten'17, Charikar-Kapralov-Nouri-Siminelakis'20). On the lower bound side, our main technical contribution is a new technique for using the INDEX problem with a large amount of side information that we hope will prove useful in other high dimensional geometric estimation problems.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。