arXiv:2502.07861cs.LGcs.AI2025-02NeurIPS被引 9

提出新算法平衡关键键值对,显著降低长文本生成内存占用。

Streaming Attention Approximation via Discrepancy Theory

论文配图:Streaming Attention Approximation via Discrepancy Theory
图 1 · 摘自论文原文
  • 基于向量平衡理论,动态筛选关键键值对以近似注意力计算。
  • 在多个长上下文任务中,比现有方法更快且更节省内存。
  • 理论证明内存下界,适合需要高效处理长序列的应用场景。

大型语言模型虽取得显著成功,但其高内存需求给长上下文的词元生成带来挑战。本文研究注意力近似在流式计算中的复杂性,提出BalanceKV算法,该算法基于几何过程选择平衡的键值对集合,依据Banaszczyk向量平衡理论实现ε-近似注意力计算。同时,我们给出了流式注意力计算的空间下界。实验表明,BalanceKV不仅在注意力近似上表现优异,在多个长上下文基准测试中也实现了端到端性能提升。

原文摘要 · Abstract (English)

Large language models (LLMs) have achieved impressive success, but their high memory requirements present challenges for long-context token generation. In this paper we study the streaming complexity of attention approximation, a key computational primitive underlying token generation. Our main contribution is BalanceKV, a streaming algorithm for $ε$-approximating attention computations based on geometric process for selecting a balanced collection of Key and Value tokens as per Banaszczyk's vector balancing theory. We complement our algorithm with space lower bounds for streaming attention computation. Besides strong theoretical guarantees, BalanceKV exhibits empirically validated performance improvements over existing methods, both for attention approximation and end-to-end performance on various long context benchmarks.

注意力机制流式计算长文本生成算法优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。