Title: Streaming Attention Approximation via Discrepancy Theory
Venue: NeurIPS 2025
Accepted: yes
Keywords: Attention approximation

Long-context Attention

Large Language Models

Discrepancy theory

Self-Balancing Walk
Abstract: 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 $\epsilon$-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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a theoretically principled KV cache compression mechanism based on recursively applying the self-balancing walk algorithm for vector discrepancy minimization. The method achieves some improvements in terms of efficiency and model quality when compared to previous techniques.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n1. The paper is well-written and motivated. I found the presentation flowing very well, and the text is written clearly and methodically. \n2. The paper is timely positioned and contributes to an important problem - the memory complexity of inference in attention-based autoregressive transformers.\n3. The approach is novel and backed by very nice theory. I found the application of discrepancy minimization algorithms very neat in this context. The paper does a good job of linking this important theoretical problem to this new domain and adapting it to its needs. Theoretically the contribution is adapting the existing method of discrepancy minimization to attention and not developing the theory anew, which the author carefully clarify, placing the paper in the correct context and not over-stating its contribution.\n4. The experiments performed are thorough and (for the most part) reflective of the statements the theory is trying to support. The authors experiment both with synthetic transformer layers and with real models, noting their method's performance over prior work.\n5. The lower bound contained in the appendix is a nice complement to this set of results. \n\n### Weaknesses\nThere aren't a lot of weaknesses, but I have a few items of note:\n1. I find the related work section a bit limited. I think it would be a good idea to add (possibly to the appendix) a few more related works on efficient attention computation, KV cache compression methods and theory for sublinear algorithms on attention, of which there is a lot. For example, there have been a few works on the sub-quadratic (offline) attention and hardness (eg. [1], [2] etc...), that, while not addressing the online streaming setting, would make for a comprehensive literature review. \n2. In the same vein, there is also a recent work [3] that focuses on lower bounds for KV cache compression. Their results appear complimentary to yours - specifically their lower bounds study the dependence on the stream length, where as yours study the dependency on the norms of the input vectors. It would be nice to discuss this briefly also in the related works section. \n3. For the experiments where you tested your method in \"real\" LLMs, what layers did you apply compression to? Did you notice a significant deterioration in model quality when the number of layers increases? My intuition says that as the number of layers increase your error will accumulate, so even if one layer is approximated well the net effect will be model degradation. Have you studied how this phenomenon appears?\n4. I'm not sure about the benefits of comparing with uniform sampling. Uniform sampling seems like it would work only in specific settings. Did you try comparing with Subgen [4]? It seems that this is another algorithm with theoretical guarantees for streaming attention.\n5. There seem to be a handful of assumptions involving the norms of the embeddings involved. Have you checked if those are generally satisfied in practice? Some comments on this would be important because the algorithm's memory depends exponentially on the square of the spectral radius $r$. I do believe that that quantity should, in general, be small, but it takes a bit away from the purity of the result. Of course, the lower bound shows the complexity is tight in that aspect, so this is not a weakness per se. Also, the error is also expressed in terms of the Frobenius norm of the value embedding matrix, which could, reasonably be large.\n6. I think a more extensive discussion of limitations and open questions left behind by the findings of this work would be a good addition to the paper overall. Currently such a discussion is very limited. \n\n[1] Alman, Josh, and Hantao Yu. \"Fundamental Limitations on Subquadratic Alternatives to Transformers.\" arXiv preprint arXiv:2410.04271 (2024).\n\n[2] Alman, Josh, and Zhao Song. \"Fast attention requires bounded entries.\" Advances in Neural Information Processing Systems 36 (2023): 63117-63135.\n\n[3] Haris, Themistoklis, and Krzysztof Onak. \"Compression barriers for autoregressive transformers.\" arXiv preprint arXiv:2502.15955 (2025).\n\n[4] Zandieh, Amir, et al. \"Subgen: Token generation in sublinear time and memory.\" arXiv preprint arXiv:2402.06082 (2024).\n\nQUESTIONS:\n1. I have not heard of the term \"streaming complexity\" before? I suppose it refers to space complexity?\n2. Line 238 typo: \"It's as\"\n3. In Theorem 3.4, the guarantee is \"for any vector\" $q \\in \\mathbb{R}^d$. However, we are only concerned with the embeddings $q_1,...,q_n$. Is there any room for improvement in the context of the self-balancing walk algorithm if you have this relaxation?\n4. What if the number of generated tokens is not known in advance? How would you set $T$?\n5. When calculating the denominator to the softmax, you are providing a vector of all $1$-s to the Balancing algorithm. That value vector has $\\ell_2$ norm equal to $\\sqrt{s}$. Is that accounted for in the analysis? \n6. In the proof of Theorem 3.4, the argument is that the self-balancing algorithm can be performed on the infinite-dimensional vectors generated by applying the exponential kernel $\\phi$. That I believe is correct, but I think some care is needed about arguing that the radius of these vectors is, in some way, bounded in terms of $r$. This seems to be a prerequisite to correctly applying Theorem A.1, and I didn't see an argument for it anywhere. Is it obvious perhaps?\n7. What is the \"prefill\" stage? It is not clear what this stage is architecturally in the experiments you are mentioning.\n\nLIMITATIONS:\nYes, though I believe that a big more extensive discussion on limitations and open questions would be a good addition to the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-21"}

comments: SUMMARY:
In this paper, the authors present a new approach to compress the KV cache in LLMs. This addresses the memory bottleneck in long context token generation. This novel technique is based on a sampling process from Banaszczyk's vector balancing theory. The authors use discrepancy theory to intelligently select which key-value pairs to retain in the cache rather than using simple uniform sampling or heuristic-based approaches.

STRENGTHS AND WEAKNESSES:
The strengths of this paper can be summarized as follows:

First, this paper has a very rigorous theoretical analysis for both the upper and lower bounds. As far as I can see, the theoretical analyses are all correct. Using Banaszczyk's vector balancing theory to attention approximation is very creative, and this is well-motivated.

Second, the experimental results seem convincing. It covers multiple tasks, such as single-layer approximation, end-to-end performance, needle-in-a-haystack tasks, and system efficiency metrics. This makes the result very influential.

The weaknesses of this paper can be summarized as follows:

First, the theoretical result of this paper requires the assumption that the norms of the query and keys are bounded. This might not hold in all the practical situations. Can the authors briefly discuss when this will hold? 

Second, it seems that this method requires tuning multiple hyperparameters, like batch size and compression rate, as shown in Theorem 3.1. It would be helpful if the authors could provide clear guidance on hyperparameter selection for different scenarios.

QUESTIONS:
Please see weaknesses

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-23

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe use of KV cache in autoregressive transformers requires storing all the keys and values, which inevitably takes $O(nd)$ memory, where $n$ is the context length and $d$ is the embedding dimension. This paper uses techniques in streaming algorithms and discrepancy theory to design an algorithm that only uses approximately $O(d^{1.5}\\exp(2r^2/\\sqrt{d}))$ memory (suppressing $\\log n$ factors), where $r$ is an upper bound on the $\\ell_2$ norm of queries and keys. An almost matching lower bound is proved additionally. The theory is further verified with experiments where they compare their techniques with previous ones.\n\nThe main technical ideas are:\n1. For each query $q_i$, we want to partition $[i]$ into two sets such that the weighted sum of the values are close, and recurse on the smaller set. Based on the self-balancing walk algorithm (which was designed for the inner product function), this paper develops an algorithm for this problem with respect to $\\exp(\\langle k,\\cdot\\rangle/\\sqrt{d})$ function.\n2. Furthermore, this paper deals with the streaming setup by partitioning the stream of tokens into batches and perform the algorithm above in each batch (merge and reduce).\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: \n1. Memory is a major bottleneck for transformers, so studying more efficient ways to store keys, queries and values is an important research area for transformers.\n2. The problem setup, high-level ideas are well explained and written, especially section 2. \n3. The empirical results solidify the theory results and illustrate that BalancedKV outperforms previous techniques on several benchmark datasets.\n\nWeaknesses:\n1. The techniques presented, self-balancing walk algorithm and merge and reduce, have both been popular techniques that are widely used. There seems to be less novelty because of this.\n2. What is a good estimation of $r$ in practice? A large $r$ will weaken the results.\n3. This paper gives error upper bound for a single attention layer, but the error could compound when we have more heads and layers?\n\nQUESTIONS:\n1. I don’t quite understand line 186 - was the previous SOTA algorithm using around $d/\\varepsilon^2$ memory? Or is the trivial $nd$ memory the previous SOTA? I am not sure whether improving from $d/\\varepsilon^2$ to $d^{1.5}/\\varepsilon$ is a big improvement.\n2. When the recursive algorithm (softmaxbalance) is performed, is the error also accumulating exponentially? Eventually the error will depend on the error between the two sets in each partition.\n3. I did not check the appendix, but is the algorithm overall easy to implement? It seems like there will be a lot of extra bookkeeping that might cause more memory usage for larger models.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces BalanceKV, a new streaming algorithm for approximating attention in large language models (LLMs) to enable efficient long-context generation. Grounded in discrepancy theory and Banaszczyk’s vector balancing, BalanceKV leverages the geometric properties of key and value tokens to select a balanced subset, achieving $\\epsilon$-approximate attention computation with reduced memory overhead. The algorithm provides provable theoretical guarantees on space complexity and approximation error under bounded $l_2$ norm assumptions, alongside a lower bound for streaming attention computation. Empirical evaluations demonstrate BalanceKV’s superior performance over existing methods like SnapKV and PyramidKV on single-layer attention approximation and end-to-end generation tasks across long-context benchmarks, including LongBench and Needle-in-a-Haystack, achieving improved memory efficiency and generation quality.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1. This paper applies discrepancy theory and Banaszczyk's vector balance to streaming attention approximation. This geometric balance method has provable approximation guarantees, in contrast to heuristic methods, and establishes a new theoretical benchmark by providing lower bounds on space complexity.\n\n2. This paper provides formal proofs of algorithmic complexity and approximation error bounds (Theorems 3.1 and 3.2), as well as comprehensive empirical evaluations. It consistently outperforms several baselines such as StreamingLLM, SnapKV, and PyramidKV on LongBench, Needle-in-a-Haystack, and single-layer attention tasks.\n\nWeaknesses:\n\n1. The theoretical guarantees (e.g., Theorems 3.1 and 3.4) rely on assumptions like bounded $l_2$ norms for query and key embeddings, which may not always hold in real-world LLM scenarios. The paper lacks discussion or empirical validation on how performance degrades if these assumptions are violated, limiting the understanding of BalanceKV’s robustness across diverse models.\n\n2. The use of a fixed compression rate (0.25) in end-to-end evaluations simplifies experimentation but may not be optimal, as different layers or attention heads might benefit from adaptive compression strategies. Exploring variable compression rates could further enhance performance or memory efficiency but is not addressed.\n\nQUESTIONS:\n1. Theorems 3.1 and 3.4 assume that the $l_2$ norms of query and key vectors are bounded. How are these norms distributed in real LLMs? How does BalanceKV perform if they occasionally exceed the assumed bounds? Can dynamically adjusting the bounds $r$ or layered processing mitigate potential performance degradation?\n\n2. Although the experiments cover context lengths from 16k to 100k tokens, how does BalanceKV perform in multimodal tasks such as ultra-long contexts (e.g., 1M tokens) or video generation? Are there other challenges in these scenarios?\n\n3. In the system efficiency evaluation, BalanceKV has the lowest Prefill Time, but the Decoding Time is slightly higher than StreamingLLM. Can you provide in-depth analysis of the specific reasons for the difference in decoding time? Is there a way to reduce the decoding latency so that it is also optimal in the decoding stage?\n\nLIMITATIONS:\nthere is no \"limitations\" section.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces BALANCEKV, a novel streaming algorithm for approximating attention in Transformer-based language models by leveraging tools from discrepancy theory. The key insight is to select a geometrically balanced subset of key-value pairs from the attention cache, thereby reducing memory usage while maintaining high approximation fidelity. The approach uses a variant of the self-balancing walk algorithm to create balanced subsets that preserve softmax-weighted value aggregations. It includes theoretical guarantees for approximation error and memory complexity, and matches known lower bounds up to polylogarithmic factors. The empirical results show consistent improvements over prior methods (e.g., PyramidKV, SnapKV, StreamingLLM) in both approximation accuracy and end-to-end performance across multiple long-context benchmarks including LongBench and Needle-In-A-Haystack.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n\n1. The paper makes a strong theoretical contribution by connecting streaming attention approximation with vector balancing in discrepancy theory.\n\n2. The algorithm is online and does not require access to future tokens, making it practical for real-time generation.\n\n3. Experiments include comprehensive ablations, runtime vs. error trade-offs, and system-level latency analyses.\n\n4. The method outperforms existing approaches like SnapKV, PyramidKV, and StreamingLLM across various setups and compression rates.\n\n### Weaknesses:\n\n1. The approach performs best when attention weights are relatively smooth; highly peaked attention distributions may suffer more from subsampling, a limitation shared by other pruning methods.\n\n2. Although the paper proves that BALANCEKV is near-optimal in a regime where ε is small and r² ≪ d, it lacks empirical analysis in extremely low-error or very long-context regimes (e.g., >128k tokens).\n\n3. The theoretical guarantees hinge on assumptions that the L2 norms of keys, queries, and values are bounded. Real-world embeddings may violate these, possibly affecting performance.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}