Streaming Attention Approximation via Discrepancy Theory
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2025
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866908907514888192 |
|---|---|
| author | Kochetkova, Ekaterina Sheth, Kshiteej Han, Insu Zandieh, Amir Kapralov, Michael |
| author_facet | Kochetkova, Ekaterina Sheth, Kshiteej Han, Insu Zandieh, Amir Kapralov, Michael |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2502_07861 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Streaming Attention Approximation via Discrepancy Theory Kochetkova, Ekaterina Sheth, Kshiteej Han, Insu Zandieh, Amir Kapralov, Michael Machine Learning Artificial Intelligence Data Structures and Algorithms 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. |
| title | Streaming Attention Approximation via Discrepancy Theory |
| topic | Machine Learning Artificial Intelligence Data Structures and Algorithms |
| url | https://arxiv.org/abs/2502.07861 |