Research map / Architectures and efficiency / Attention and transformers
Streaming Attention Approximation via Discrepancy Theory
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes BalanceKV, a streaming KV-cache compression algorithm built on Banaszczyk-style vector balancing (SoftmaxBalance + merge-and-reduce) that provably approximates attention with sublinear memory, and shows a matching-style space lower bound.
Problem
Approximating softmax attention in a streaming setting with sublinear memory, so that the key-value cache of LLMs can be compressed during long-context token generation.
Approach
Extends the online self-balancing walk from linear vector balancing to the exponential kernel exp(<k,q>/sqrt(d)) via a feature map, runs it recursively in a merge-and-reduce tree, and buckets tokens by value norm to get a streaming estimator of attention with memory roughly d*sqrt(d)/eps instead of 1/eps^2 for independent sampling.
Key results
- On LongBench-E with Qwen2.5-14B-Instruct at roughly 0.25 compression, BalanceKV averages 50.62 versus 50.22 for SnapKV and 49.88 for uniform sampling.
- On Needle-In-A-Haystack with Llama-3.1-8B-Instruct at roughly 0.25 compression, BalanceKV reaches 0.99 average accuracy versus 0.83 for SnapKV and 0.90 for uniform sampling.
- Theory: BalanceKV uses about d*sqrt(d)/eps memory versus 1/eps^2 for independent subsampling, and a space lower bound of Omega(min{1/eps^2, d exp(2r^2/sqrt(d))}) is proven.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use BalanceKV instead of SnapKV | Qwen2.5-32B-Instruct cache compression · LongBench-E | average: 48.77 -> 48.84 | Helped |
| Use BalanceKV instead of SnapKV | Qwen2.5-14B-Instruct cache compression · LongBench-E | average: 50.22 -> 50.62 | Helped |
| Use BalanceKV instead of Uniform | Llama-3.1-8B-Instruct cache compression · LongBench-E | average: 46.38 -> 46.77 | Helped |
| Preserve strongly anti-correlated-key tokens before applying BalanceKV to remaining tokens | Llama-3.1-8B-Instruct cache compression · Needle-In-A-Haystack | average accuracy: 0.83 -> 0.99 | Helped |
| Use BalanceKV augmented instead of PyramidKV | Llama-3.1-8B-Instruct cache compression · Needle-In-A-Haystack | average accuracy: 0.90 -> 0.99 | Helped |
| Use BalanceKV augmented instead of StreamingLLM | Llama-3.1-8B-Instruct cache compression · Needle-In-A-Haystack | average accuracy: 0.31 -> 0.99 | Helped |
| Use BalanceKV augmented instead of Uniform | Llama-3.1-8B-Instruct cache compression · Needle-In-A-Haystack | average accuracy: 0.90 -> 0.99 | Helped |
| Increase block size (ablation) | BalanceKV single-layer attention approximation · TriviaQA, Llama-3.1-8B-Instruct layers 1 and 15 | relative error | Helped |
| Increase block size (ablation) | BalanceKV single-layer attention approximation · TriviaQA, Llama-3.1-8B-Instruct layers 1 and 15 | runtime | Hurt |
Limitations
- Theoretical guarantees assume bounded query and key norms; the paper validates this empirically only on TriviaQA with Llama-3.1-8B.
- The Needle-In-A-Haystack result uses an augmented variant that deterministically preserves anti-correlated tokens, not the pure algorithm.
- Runtime per step scales as d^2/eps^2, and the prefill overhead exceeds the uncompressed baseline.
Experimental setup
Datasets: LongBench TriviaQA Needle in a Haystack MS COCO
Benchmarks: LongBench Needle in a Haystack TriviaQA MS COCO
Metrics: relative attention approximation error average accuracy wall-clock latency
Baselines: StreamingLLM SnapKV PyramidKV uniform sampling ClusterGen
Builds on: self-balancing walk merge-and-reduce SnapKV
Closest papers by meaning
- Efficient Streaming Language Models with Attention Sinks (ICLR 2024 · similarity 0.90)
- ArkVale: Efficient Generative LLM Inference with Recallable Key-Value Eviction (NeurIPS 2024 · similarity 0.86)
- MagicPIG: LSH Sampling for Efficient LLM Generation (ICLR 2025 · similarity 0.86)
- CateKV: On Sequential Consistency for Long-Context LLM Inference Acceleration (ICML 2025 · similarity 0.86)
- Unveiling and Harnessing Hidden Attention Sinks: Enhancing Large Language Models without Training through Attention Calibration (ICML 2024 · similarity 0.86)
- Recurrent Attention-based Token Selection for Efficient Streaming Video-LLMs (NeurIPS 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.