Streaming Attention Approximation via Discrepancy Theory

Thank you for carefully studying our submission and the additional, very useful suggestions and comments!

-->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. 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…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.

Response: Thanks a lot for pointing out this as well as related work on efficient attention and kv cache compression, and lower bounds. We will incorporate these and add relevant discussion of the limitations of this work and future directions in the final version of the paper.

-->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?

Response: In our end to end experiments we apply the compression to all layers. Regarding the latter part of your question, we perform the following experiment - we evaluate the end to end performance of Llama 3.1 8B Instruct after applying the compression on only the first i layers on Qasper (one of the datasets in LongBench), for the values i in [0,8,16,24,32] (32 is the last layer), and observe the scores of [42.8,42.4,39.8,.38.2,35.8] respectively. This confirms your intuition that as the number of layers on which the compression is applied increases, the performance decreases.

-->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.

Response: The fundamental goal of the paper is to demonstrate the power of the geometric sampling in streaming attention approximation. Therefore, geometric sampling is compared to uniform sampling as baseline.

We further benchmark Subgen [4] and compare it to our BalanceKV using Llama-3.1 model and a subset of LongBench datasets (due to limited timeline). The experiments were performed with a single NVIDIA RTX A6000 GPU with 48GB VRAM. The below table summarizes the performance of each method on datasets we test.

As a result, our BalanceKV outperforms Subgen on all datasets but “multifieldqa_en”. We will add complete experiments in our final draft.

-->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 . 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.

Response: Thanks for the question. Please refer to point 2 of the response to Reviewer BB75 for the discussion of the upper bounds on the norms of the embedding vectors. Regarding the choice of the Frobenius norm - this is a great suggestion and it is a very interesting open problem to obtain the same space bounds as ours, but with error guarantee scaling with the spectral norm of V. Note that the inequality $||V||_F \leq \sqrt{d}||V||_2$implies that by rescaling the epsilon we can use BalanceKV to get an algorithm with the standard error of attention approximation, whose space and time complexity is scaled by $\sqrt{d}$ and $d$ respectively. Furthermore, we observe in practice that the singular values of the matrix V decay fast, thus the theoretical upper bound of $\sqrt{d}||V||_2$ on $||V||_2$ is loose.

Questions:

1. I have not heard of the term "streaming complexity" before? I suppose it refers to space complexity?

Response: You are absolutely correct, we refer to the space complexity of a streaming algorithm solving the problem of attention approximation when we write ``streaming complexity of attention approximation’’.

2. In Theorem 3.4, the guarantee is "for any vector". However, we are only concerned with the embeddings $q_1, \ldots, q_n$. Is there any room for improvement in the context of the self-balancing walk algorithm if you have this relaxation?

Response: Indeed, the query embeddings definitely have special properties; this is evidenced, for example, by the existence of attention sinks (https://arxiv.org/pdf/2309.17453). In more detail, it has been observed that the initial input tokens have large attention scores, which, in particular, suggests that the query embeddings must be non-trivially correlated with the key embeddings of the initial input tokens. It is an exciting direction for future work to discover more geometric properties of the query embeddings and leverage them to improve BalanceKV.

3. What if the number of generated tokens is not known in advance? How would you set T?

Response: In practice, the maximum number of generated tokens is a property of the model and, therefore, it is known in advance.

4. 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?

Response: When we compute the approximation to the denominator we set all of the value vectors trivially equal to scalar 1. In other words, they are indeed vectors of all 1-s but the dimension, s, is 1. To briefly justify this choice of value vectors, note that Theorem 3.4 proves theoretical guarantees of SoftmaxBalance which may be applied to approximate $\sum \exp(\langle k, q\rangle/\sqrt{d})v$ where the dimension of the value vectors is arbitrary. Since both the numerator and the denominator of the attention function may be expressed in this way, we approximate both with the same algorithm, just applied to different sets of value vectors. Setting all value vectors to be scalars 1, the expression $\sum \exp(\langle k, q\rangle/\sqrt{d})v$ becomes precisely the denominator of the attention function.

5. 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?

Response: Indeed, we apply Theorem A.1 to vectors $\varphi(k)$ and $\varphi(q)$, and the guarantees of Theorem A.1 depend on the $\ell_2$ norms of these vectors. In the expression in between lines 918 and 919 we note that $||\varphi(k)||^2_2 = \exp(||k||^2_2/\sqrt{d})$. Ultimately, we use our assumption $||k||_2, ||q||_2 \leq r$, which gives the desired upper bound on the norms of the infinite dimensional vectors. Note, however, that Theorem A.1 is only used in the proof of Theorem 3.4, which does not assume anything about boundedness of norms of key and query vectors and whose guarantees are stated in terms of max $||k||^2_2$. This is why we did not mention r in that proof. We will add a discussion regarding this in the final version, to improve the understanding of the big picture.

6. What is the "prefill" stage? It is not clear what this stage is architecturally in the experiments you are mentioning.

Response: ''Prefill'' refers to the initial forward pass over the entire input prompt to populate the model’s KV (key-value) cache. During the next stage, referred to as ''generation'' or ''decoding'', the new tokens are generated one by one, and in this process it is crucial to compute the attention function between query embeddings of each newly generated token and key value embeddings of each past token, which implies the need to keep the entire KV cache during ''decoding''. In our experiments, we do not optimize the ''prefill'' stage and compute the exact KV cache with standard methods. However, during the ''prefill'' stage we run the BalanceKV algorithm to compress the exact KV cache at every head and every layer to decrease the space complexity during the ``decoding’’ stage.

Thank you for carefully studying our submission and the additional, very useful suggestions and comments!

1. 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?

Response: Thanks for the question, please refer to point 2 of the response to Reviewer BB75 where this is discussed.

2. 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.

Response: The compression rate is, in practice, dictated by the memory budget. As we discuss in the response to question 1 of Reviewer pBkK, there is only one hyperparameter in our experiments which requires tuning - the parameter gamma which is a function of the upper bound on the norms of key and query vectors. We observed that quite a wide range of choices for this parameter lead to similar performance in practice, and we fix it to be 4 across all experiments.

Thank you for carefully studying our submission and the additional, very useful suggestions and comments!

-->What is a good estimation of r in practice? A large r will weaken the results.

Response: Thank you for pointing this out. Please refer to the 2nd point of the response to Reviewer BB75 for the discussion of the upper bounds on the norms of keys and queries in the real data.

-->This paper gives error upper bound for a single attention layer, but the error could compound when we have more heads and layers?

Response: Thank you for this observation! Indeed, the error of BalanceKV could compound across multiple layers. We show this does not lead to significant performance degradation in practice - in our end-to-end experiments we show applying BalanceKV across all layers and heads leads to better performance as compared to other state of the art methods.

Questions:

1. I don’t quite understand line 186 - was the previous SOTA algorithm using $d/\epsilon^2$ memory? Or is the trivial memory nd the previous SOTA? I am not sure whether improving from $d/\epsilon^2$ to $d^{1.5}/\epsilon$ is a big improvement.

Response: Thank you for your question! Indeed, uniform sampling has the theoretical guarantee of $d/\epsilon^2$ and is the previous SOTA. From the theoretical perspective, BalanceKV outperforms uniform sampling in the low error regime: more precisely, when $\epsilon < 1/d$. However, as our experiments show, in practice BalanceKV outperforms uniform sampling for all ranges of $\epsilon$. This suggests that instances arising in practice do have geometric structure favorable to our discrepancy-based BalanceKV algorithm. This holds true for both attention approximation and end-to-end experiments.

2. 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.

Response: Thank you for your question! The error of MergeAndReduce - the recursive implementation of SoftmaxBalance, depends linearly on the depth of the recursion $T$. We formally prove the theoretical guarantees of MergeAndReduce in Subsection A.3.2, but briefly this may be explained as follows. MergeAndReduce is an online implementation of the following procedure: we split our dataset into blocks of size $t$, compress each block by a factor of 2 with the SoftmaxBalance algorithm and then recurse on the resulting smaller dataset. Please refer to Figure 2 in the paper for visual illustration. The error of transitioning from the $i-1$-st to $i$-th recursion level is upper bounded by the product of $2^i$, the number of blocks of size $t$ into which we partition the dataset at the $i$-th level of recursion, and the error of SoftmaxBalance on one block of size $t$. Note that the number of blocks in the partition of the dataset at the $i$-th level of recursion is upper bounded by $n/(2^i \cdot t)$. Therefore, our upper bound on the error at the $i$-th level of recursion is independent of $i$.

3. 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.

Response: Thank you for your question! In a nutshell, the algorithm is easy to implement. More precisely, in our experiments, we use BalanceKV to compress the KV cache at the prefill stage. In the streaming part, we generate only a few tokens, therefore compression is not necessary. As we discuss in more detail in point 1 in the Questions section in the response to Reviewer pBkK, there is a single parameter that we need to tune, and it is quite easy to choose this parameter empirically.

Thank you for carefully studying our submission and the additional, very useful suggestions and comments!

--> The theoretical guarantees (e.g., Theorems 3.1 and 3.4) rely on assumptions like bounded 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.

Response: Thanks for pointing this out. Please refer to the 2nd point of response to Reviewer BB75 for the discussion of the upper bounds on the norms of keys and queries in the real data. In practice, BalanceKV is robust to the choice of assumed upper bounds on the norms of key and query embeddings. We discuss this more thoroughly in the response to your question 1.

-->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.

Response: Thank you for this interesting suggestion! We certainly agree that analysing the properties of different layers, attention heads and designing an adaptive compression scheme based on their importance is an exciting future work direction.

Questions:

1. Theorems 3.1 and 3.4 assume that the 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 or layered processing mitigate potential performance degradation?

Response: Thank you for the question. For the distribution of norms, please refer to the 2nd point of response to Reviewer BB75. BalanceKV is empirically very stable. In particular, while norm upper bounds are indeed needed for the theoretical analysis of BalanceKV, its practical performance is stronger. In the practical implementation of our algorithm the expression 1/(2c*R^2), which appears in line 8, is a single hyperparameter, parameter gamma in our implementation. This is the only hyperparameter that our implementation needs and it is the only parameter which depends on the norms of key and query vectors. We did an ablation study on the parameter gamma, which showed that the range of values of gammas which allows for good empirical results for BalanceKV is quite wide, suggesting strong geometric structure. We will include the ablation study in the final version of the paper. In our experiments, we set gamma to a small constant (about 4).
Studying adaptive variations of BalanceKV is a great direction; however, as practice suggests, the norms of key and query vectors concentrate well, and imprecise choice of the assumed upper bounds does not lead to performance degradation. Therefore, dynamically updating the assumed upper bound on the norms of embeddings is unlikely to be helpful.

2. 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?

Response: Thanks for the question. We perform the evaluation of BalanceKV and uniform sampling when applied to the InternVL2.5-8B multimodal LLM for compression rates 1/4 and 1/16, for evaluation on the MS COCO image captioning dataset, the scores are reported below. (The bracket for every method contains the compression rate). The experiment was run on 4 X A100 GPU each with 80 GB VRAM.

For example, for the CIDEr score (specifically designed for image captioning and evaluates how well a generated caption aligns with a set of human references, higher is better), our method is consistently better than the baseline. The ultra long context or video generation evaluation is also a great suggestion, we will include it in future work.

3. 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?

Response: Thank you for the insightful question. Indeed, under idealized conditions, one would expect that decoding time across compression methods to be identical, since the per-token compute in the decoding stage remains the same across methods once KV cache is compressed. However, in practice, several system-level factors can contribute to minor differences in the decoding time, e.g., slicing multi-head KV caches or cache locality. We believe that further system-level tuning (e.g., kernel optimization) could close the gap and potentially make BalanceKV optimal in both prefill and decoding stages.

Thank you for carefully studying our submission and the additional, very useful suggestions and comments!

1. 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).

Response: We think these are great suggestions. Extremely low error regime corresponds to the low compression rate regime. We conduct the experiment where we apply BalanceKV and uniform sampling to a few of datasets in LongBench with low compression rates 0.8, 0.9 and 0.95. In our evaluations, we use the model Llama-3.1-8B-Instruct. The experiments are performed on a single A100 GPU with 80GB VRAM. Our experiments show that even in this regime BalanceKV outperforms uniform sampling.

Extremely long context regime is also a very interesting direction, and we will include it in future work.

2. 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.

Response: Thank you for raising this excellent point! Indeed, the assumption on the boundedness of the embedding vectors is both important to our theoretical guarantees and stemming from the properties of the real-world datasets. In our implementation of BalanceKV, we run the self balancing walk on the datasets of keys shifted by their average. Note that the operation of shifting the key vectors causes both the numerator and the denominator of the softmax function to scale by a constant factor, therefore if our algorithm successfully approximates the attention scores for the shifted keys, then it also performs well for the original set of keys. We compute the $\ell_2$ norms of queries and keys shifted by their average, and values (qkv) using the pretrained LLaMA-3.1-8B-Instruct model and Qasper (one of the datasets in LongBench) dataset. Specifically, we analyze each prompt in TriviaQA and compute the average L2 norms of all qkv vectors across all layers and attention heads during the pre-fill stage. Due to the space limit, we provide representative results in the below table from (randomly chosen) prompts of various sequence lengths. The key findings are that all qkv norms consistently concentrate around some constants (~15 for query, and ~15 for key, ~3 for value) with small confidence intervals (CI). Importantly, the norms remain stable across a wide range of sequence lengths, suggesting that these norms do not grow with input length. The table below clearly indicates that the query/key/value norms are well-bounded and stable, validating that our assumption is not only theoretically sound but also holds in practice for real-world language model usage. We will add more detailed results in the final draft.

The intuition is that the qkv embeddings typically can be obtained by multiplying weights by the input embeddings, and those input embeddings are operated through the layer normalization. Therefore, entries of qkv embeddings are expected to be small assuming the weight matrices have small eigenvalues independent of the sequence length. So the norms of qkv embeddings in our case after this normalization are expected to be small constants regardless of the sequence length.