Streaming Attention Approximation via Discrepancy Theory
Authors: Ekaterina Kochetkova
Kshiteej Sheth
Insu Han
Amir Zandieh
Michael Kapralov
Venue: NeurIPS 2025 spotlight ()
Keywords: Attention approximation
Long-context Attention
Large Language Models
Discrepancy theory
Self-Balancing Walk
Primary area: theory
TL;DR: We design a geometric sampling framework to enhance memory efficiency and precision in long-context LLMs.

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.

Streaming Attention Approximation via Discrepancy Theory
Ekaterina Kochetkova
Kshiteej Sheth EPFL
EPFL
okaterina. how.brrnaa.kchhckkoooeeee....... ch
kshiteej shethhepfl.cch
Insu Han
Amir Zandieh Google Research zandi eh@googee com
Michael Kapralov EPFL
KAIST insu. hanakaist....... www.:kaiiitt.....k.........
michael. .apaaollapr.ll
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, key computational primitive underlying token generation.
Our main contribution is BalanceK a streaming algorithm for e-approximating attention computations based on geometrie process for selecting a balanced col- lection of Key and Value tokens per Banaszczyk's vector balancing theory. We complement our algorithm with space lower bounds for streaming attention com- putation. Besides strong theoretical guarantees, BalanceKV exhibits empirically validated performance improvements over existing methods, both attention approximation and end-to-end performance various long context benchmarks.
Introduction
Transformer-based models are the foundation of ongoing artificial intelligence revolution. Their applications span wide range of domains, from leading-edge language models (LLM) [1, 65] to text-to-image [58, 66, 69], text-to-video synthesis [70], coding assistance [68] and even in multimodal domains across text, audio, image, and video [53]. At the core of these models is the Transformer ar- chitecture, powered by the self-attention mechanism [73], which enables effective capture of pairwise correlations across tokens in an input sequence. As these models scale in size and context length [41], they face significant computational challenges, particularly in terms of memory usage. Efficiency and accuracy are essential unlock the full potential of L.MMs in generating long sequences.
Space bottlenecks in transformer models. Most large language models, along with multimodal and video models, adopt an autoregressive, decoder-only architecture. This architecture generates tokens sequentially, applying attention dynamically to each newly generated token. To avoid redun- dant attention
computations during the generation phase, these models explicitly store the key and value embeddings of previously generated tokens in a cache in each attention layer. Thus, a major challenge is the fact that the memory complexity of storing previously generated key value embeddings scales with both the model size (i.e., the number layers and attention heads) and, critically, the context size. Additionally, each model session typically requires its own dedicated cache for storing key value embeddings, further exacerbating memory usage. This growing demand has become significant bottleneck, affecting both memory consumption and computational speed, particularly for models handling long context lengths.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Streaming attention computation. The main reason for the need of storing the past key and value embeddings is for the attention computation happening inside each self attention layer during token generation after processing context - to generate the next token, each self attention layer computes the attention between the query embedding of the current token and the key and value embeddings of all the tokens that were previously generated or part of the context. In this paper we study the streaming attention approximation problem - the problem of approximately computing attention using small amount of space, i.e. without storing all previously seen key and value embeddings. Our main contribution is BALANCEKV, a novel provably correct algorithm for streaming attention approximation based on discrepancy theory. The core of our approach is a vector balancing algorithm from discrepancy theory that exploits the geometry of key and value tokens to deduce a small subset of them that well approximates the operations happening inside self-attention layer. We complement our algorithm with a lower bound on the streaming complexity of approximating attention.
An algorithm for streaming attention approximation can directly be used for compressing the key value cache which stores the past key value embeddings in each layer in an LLM,, thus improving the efficiency of LLM token generation. We empirically evaluate BALANCEKV both on the problem of approximating attention and on end-to-end generation tasks, showing performance gains.
1.1
Related Work
For discrepancy theory. Banaszczyk's ' seminal works [6, 7] establishing theoretical guarantees for vector set discrepancy have sparked research in the vector balancing problem [17]. This led to algorithmic developments in both offline [8] and online [9, 3, 43] settings. The vector balancing problem has particular relevance to streaming and sublinear algorithms, as minimizing dataset's discrepancy yields small subsets that effectively preserve the original dataset's properties. Recently [55, 15] extend these discrepancy theory ideas for kernel density estimation using sublinear memory.
A simple yet effective approach is quantizing previously generated key value embeddings with fewer bits [80, 78, 26, 40, 49, 35, 84, 81]. Another line of work focuses on token-level pruning, where redundant or less important tokens get evicted from the set of all previously generated key value embeddings [10, 85, 48, 76, 83, 46]. Many of the works in this line have used accumulated attention scores to select important previously generated tokens [85, 46, 76]. Recent works extend those methods to an adaptive way of budget allocation across layer [14] and head [30].
1.2 Overview of Our Contributions
In this work we take the token subset selection approach to reduce the memory complexity of LLM token generation: store and maintain only subset of previously generated key and value embeddings corresponding to few "important" tokens in the sequence. Of course, the central question is how to define "importance" of tokens. Our approach here is to apply discrepancy theory, which, at a high level, considers a token important if it is crucial to preserving the projection of the total collection of tokens onto some direction in the token space. This leads the idea of selecting a subset of tokens that is "balanced" simultaneously in every direction. Inspired by the recent breakthrough result of [3] on online discrepancy minimization, we design a method for balancing key-value pairs online using small space, namely our BALANCEKY V algorithm. Interestingly, this algorithm is online, i.e. the importance of a token is determined only by preceding tokens- in sharp contrast with state of the art heuristics for token selection such as PyramidKV [14] and SnapKV [46], whose performance, show, our algorithm matches or improves upon. Our contributions are:
1. In Section 3 we propose BALANCEK an algorithm for recursively compressing the set of previously generated tokens using a geometric correlated sampling process based on discrepancy theory. We show that BALANCEKV V gives provable guarantees for streaming attention approxi- mation under the bounded f norm assumption (Theorem 3.1). Using tools from communication complexity, we also show a lower bound on the memory complexity of any algorithm for stream- ing attention approximation in Section 3. Section 2 contains the formal problem formulation of streaming attention approximation, its applicability to key value cache compression, as well as a technical overview of the main results and techniques of Section 3.
2. In Section 4 we empirically evaluate our algorithm in various settings. In Section 4.1 we show our approach leads to a lower relative error for single layer attention approximation for open-source LLMs including Llama-3.1 1-----IIIIIt-ueee [27] and Ministral-8B I--II--rItttt

---

2410 [52] as compared to uniformly sampling keys and values in the cache. Section 4.1 we also perform ablation studies to show how various parameters in our algorithm affect the relative error for single layer attention approximation. In Sections 4.2 and 4.3 we perform end to end experiments on various benchmarks such as LongBench [5] using models of various sizes such as Llama-3.1 1--....1.-B.1. and Needle in a Haystack [39]. We show that our provable method for attention approximation when applied to key value cache compression performs better compared to previous existing token subset selection heuristics on end to end tasks. Finally in Section 4.4 we present system efficiency metries regarding our implementation.
and Qwen-2.5-322 2-222------IS 32B-Instruct [77, 71],
2 Technical Overview
In this section, we first set up the formal problem formulation that we tackle, followed by an overview of our techniques and our main results.
2.1 Streaming Attention Approximation: Formulation and Motivation
Autoregressive Transformers generate tokens one by one and each depends the previously gener- ated tokens. When Transformers process a sequence of tokens, the attention mechanism operates by computing three types of embeddings for each token every layer: query, key and value. The query and key capture how different tokens interact, while the value the actual content to be aggregated. Such interactions are quantified by so-called attention scores, obtained by applying the softmax the inner product between the query of a given token and the keys of all others. These scores determine how much each previous token's 's value contributes the final output. Once the keys and values computed for a given token, they do not need to be recomputed when generating subsequent tokens.
Formally, suppose that we have (g1, k1, v1)..... (9n, (9n. Kn, Un), that is the j-th token is represented as a triplet of (q): ky, Uy) where qj ky, Uj E R2 for j E [n]. Let K,, V3 E iRi be matrices defined by stacking those keys and values in their respective rows. To compute the following every step j generate j token, called the streaming attention problem:
stream of query, key and value embeddings
K5-9) 90 Vd
T
Attn(gj: K3, V2) = softmax
Vj.
(I)
Keeping all of the key-value pairs in the cache is prohibitively expensive, especially for long sequences. Instead, we opt for approximate computation by sampling a few key-value pairs. Specifically, our goal is to construct an algorithm that at every time step computes an estimator z; for Attn(gj, Kj, Vj) sublinear in n time and memory, In particular for given precision following error constraint:
O, 2, should satisfy the
2 Attn(g, K,, Vi)lla S softmax (K, %
Va ){
(2)
A sublinear in time and memory algorithm compute 2j will require knowledge of significantly less key-value pairs than K,. Vj. thus reducing the size of the key value cache needed to store them. This motivates the study of streaming attention approximation, used for key value cache compression during LLM token generation. In the next section we discuss how will construct such an estimator zj at high level.
an algorithm for this can directly be
2.2 SOFTMAXBALANCE. Attention Approximation via Discrepancy Theory
We now start with presenting the main ideas of our approach. By the definition of softmax, Equa- tion (1) can be written as
K2 Vd
qj
T V2,
Attn q3 Ky, =
Z, exp
where for a matrix A we write exp(A) A) to denote entry-wise exponential function to A and Z, == Σiegi exp( //// Our approach to approximate Attn(y;, K5, Vý) consists of two subroutines which approximate:
3

---

1. Softmax normalization 2, = ExeDi exp( 9/) Vd),
2. Matrix-vector product between VÚ and exp(K) gj/vd).
To understand our main idea, suppose we are at the end of the stream j= and we store all key-value pairs (ki, vi)....... (Ki, Un). Then for arbitrary query 4n we aim to approximate the matrix-vector product exp( Ka qa/ /v)"" .V.. Eiejn exp((ki, gnn/v/dl Vd)u, by choosing subset of compression rate of 0.5. Suppose
the rows of K, and Vn of size at most n//2 which corresponds we can design algorithm which splits the set C' of all keys and values into two groups C" and CNCC so that the matrix-vectis vector product function for any query vector qn. roughly equal over and C2 that is informally,
(k, (n) Vd
(kc, (n) (v/za)
Σ ECC
Σ exp
exp
Vd
(K,))C
Then, we are able to approximate the matrix-vector product function with either one of the sums above since informally:
(ki (k. qm) Vd
(k, (k) Gn. Vd
Σ EK
Σ exp (A,
exp
1 =
Therefore, it would suffice to keep the smaller subset of C' and CN,C" as the desired subset of key value embeddings and discard the rest. If we wanted to compress the key value cache to smaller size by a factor 2n for some T, we would recursively compress the selected subset using the same procedure T more times.
similar goal is captured by the vector balancing problem studied extensively in discrepancy theory; given set of vectors C . knn C IRd with [ S for all j. partition them into two groups C", C \ C' such that for any E Fd it holds Σxeo (k, q) ~ Σkecro (k, q) with high probability. The Self- half-Balnncing Walk algorithm [3] is breakthrough result for the above vector balancing problem. However we need to develop an algorithm for the vector balancing problem with respect function exp((k,)) +) /dyo instead of the inner product function (k. -).
Our first contribution is to develop an algorithm for our task, building upon the result from the self-balancing walk [3], which essentially randomly partitions the of keys and values C' into C' and C \\ C' such that the following holds with high probability under the assumptions that the norms of the query and key embeddings are bounded,
Σ exp
(k, (n) Vd
Σ (K.UREC
exp
(k, 9m) Vd
< O (log(nd)) max ||//]|2- jE(n]
(R.O)EC'
We refer to this algorithm as SOFTMAXBALANCE its formal guarantee is presented in Theo- rem 3.3 and its pseudocode is presented in Algorithm 2. Theorem 3.3 shows that SOFTMAXBAL- ANCE succeeds to divide C' into subsets C' and C"(C" which are balanced with respect to function exp((k, -/V/do up to an error which only has logarithmic dependence on the size of C. In addition, can accept as input value vectors of arbitrary dimension 8. Therefore, if instead 1, we will get an
SOFTMAX
of the value vectors .,...., algorithm for the vector balancing problem with respect to function exp( (ks /)/ /V).). This implies that we can use SOFTMAXBALANCE to compress the key value cache to even approximate the soft- max normalization Σien exp( (k,, qu)) /d). We now discuss how to use SOFTMAXBALANE streaming attention approximation, i.c. to use it to compute an estimator satisfying Equation (2).
Rd we input the set of scalars 0
for
2.3 BALANCEKY: Implementing SOFTMAXBALANCE Streaming
For sequence of n tokens and given memory budget of < n. we aim to design procedure which applies SOFTMAXBALANCE to select from n key-value embeddings set at most in the streaming setting and can compute an estimator zj satisfying Equation (2) for all steps in the stream. the streaming setting one needs to consider the following aspects. As described in the previous section, one iteration of SOFTMAXBALANCE only allows one to select nz/2 sized subset of

---

n key-value embeddings, which is higher than the desired budget of embeddings. This can be easily mitigated by recursively applying SOFTMAXBALANS key-value embeddings. However, this cannot be implemented in the streaming as we have limited memory budget of which prohibits us from storing all key-value embeddings during recursion.
2log(n/4) times, each time halving the set of
To deal with this,
the classical merge and reduce technique used in the design of streaming algorithms [13, 51, 33J. MERGEA NDREDUCE algorithm is a recursive binary tree-based approach recursively in a streaming setting with the total memory not exceeding Ödt), where Ō(-) supresses polynomial in log n factors, under the assumption that the norms of queries and keys are bounded. The guarantees of MERGEANDREDUCE are presented in Theorem 3.4, its pseudocode in Algorithm 4 and a visual representation in Figure 2. If the norms of all value embeddings in the stream are the same up to constant factors, that is for all i E [n] 0.5 < ||o/|l/|||v, I 2. then the outputs of MERGEÅnDREDUCE can be used to construct an estimator 2j satisfying our attention approximation guarantee of equation Equation (2) with precision E for f Ö(VdfE). However, the value embeddings may have very different norms.
that allows one to implement SOFTMAX
Our main algorithm BALANCEKV (pseudocode in Algorithm 1) deals with this issue by grouping the key-value embeddings in the stream according to the norms of the value embeddings, running separate instance of MERGEANDREDUCE on each group, and combining the outputs each instance of MERGEANDREDUCE BALANCEKV constructs a final estimator zj satisfying Equation (2) with precision E only using ((dvd/e) memory and (dd /=a) runtime per every step j of the stream, assuming the norms of query and key embeddings are bounded. Existing methods [83] subsample keys and values independently in the cache, and thus have a 1/e² dependence on E in total memory. The guarantees of BALANCEKV V presented in Theorem 3.1.
Finally using the lower bound on the communication complexity of INDEX, we show a lower bound on the memory complexity of any algorithm for streaming attention approximation in Theorem 3.2.
3 Main Theoretical Results
Our main algorithm for streaming attention approximation is BALANCEKV, It takes in as input kn, vn)) and at every step of the stream outputs an estimate 23 to Attn (gy. yj Vj) (see Equation (1) for the definition of Att(()) satisfying Equation (2) with precision e. Assuming that the &e norms of qj. k2 are at most r for all j. BALANCEKV uses
stream of n tokens (92, k1, v2), (42, k2, Uy)..... (4n,
total space ((dd/der//va /va 1/E) and uses e (r22a 1/E2) runtime at each step j of the stream to output zj. Our main theorem is as follows.
O, any positive integers n,d, any set of tokens (g), ki, v1), (42, k2, u2)....; (qn, kny vn) where qj, kj, E Rdd satisfy 1|0 [kllla S r for all consider an invocution of BALANCEKV with
Theorem 3.1. For any T,0
batch size t t= Ō 1====0((22/@/ (Vde2P/Va1)) and compression rate 2 with T log(n/t).
Then BALANCEKV outputs a vector 2j satisfying Equation (2) with probability at least 1 - 1-1/poly(n) at every step j of the stream. It uses lotal memory Ō (dvde>> (V across all steps of the stream and runtime of o (dPArr2/V
per step of the stream.
A pseudocode of BALANCEK V is described in Algorithm 1. At its core BALANCEKY relies on our main discrepancy based algorithm, namely SOFTMAXBALANCE- see Section 3.1 for details on SOFTMAXBALANCE. BALANCEKV uses the output of SOFTMAXBALANG to compute estimates of the numerator and denominator of Attn(g, KG V2) and returns the desired attention approximation zj for each streamed index j. There are two subtleties, however. First, it is important to bucket tokens in the stream according to the norm of the value vectors see lines 5 and Second, direct application of SOFTMAXBALANCE would require too much memory space. To ensure small space usage, we apply classical streaming technique, namely the MERGEANDREDIE top of SOFTMAXBALANCE to reduce the space consumption. The space reduction achieved by MERGEANDREDUCE by running logarithmic number of copies of SOFTMAXBAL.N ANCE in tree-like fashion. More details are introduced in Section 3.2.
algorithm on
5

---

Algorithm 1 BALANCEKV ((g), hig, 11) T.E)
1: input: stream of rt tokens (g), Kg, 12), diameter batch size t, compression rate 2 precision parameter E.
2: // Bucket the stream and maintain log(n) instances of MERGEANDREDDUE,, MR-NUMERATOR, for each bucket approximate the numerator of Attn(g, K,, V2); and one instance, MR DENOMINATOR, to approximate its denominator.
3: Umax f-0
4: repeat
5: Find an index such that 2' 2 110y122 2 2
6: Send (ky vU) imput to MR- AR-NUMERATOR,
// Bucket the stream by |
7: Umax < max fll Umax]
8:
Erase all MR-NUMERATOR, with 2' S 2n v Umax
// Erase small buckets
9:
Oo C..... Cf the output of MRRNNMEERTOR,,
10:
V F U,C for l = ......T
// Combine the outputs of MR-NUMERATOR;
11:
Send (kj, 1) as imput to MR- AUR-DENOMINATOR
12:
K KT f .R MIR-DENOMINATOR
Ii Eik.oeev exp
13:
output: 2j
Σt 2' E CAMMEN exp
(k:0g)
14:
j 5 j
until token stream ends
To summarize, BALANCEKV groups tokens in the stream according the norms of the corresponding value embeddings. runs separate instance of MERGEANDREDUCE on each group, and combines the outputs of each instance to construct the final estimate for Attn(gj, Kj, V, at each step j [n]. Next we present SOFTMAXBALANCE and MERGEA NDRRUUFF, The full proof of Theorem 3. is given in appendix Section A.1. Finally we state the theorem which provides lower bound on the memory complexity of any algorithm for streaming attention approximation below, its full proof is provided in appendix Section
Theorem 3.2. Suppose that y2 < d. Any streaming algorithm which on input (kn "nb.g), q). [1q]|2, [|kill2 S r, outputs 2q satisfying Equation (2) with probability 0.999 has space complexity Q (mint (exp(2r//Vd))
3.1 SOFTMAXBALANCE
We now present our main discrepancy based compression algorithm, SOFTMAXBALANCE. Given sequence of key and value embeddings C {(k1, 01), vn)} (with key and value embeddings having possibly different dimensions), the goal of SOFTMAXBALANCE is to pro- duce partition of C into subsets C²,C C" such that for any query & Rd we have that Exke)eer exp (ks 4)/ /d) E(K,))Er exp((k, q) Vdju with high probability. Without loss of generality assume that |C'| S CI/2, we can then output Eik exp (k:, 4// //// approximation to (( exp((k, 9) / /d)e,, thus achieving factor 2 compression. Its description presented in Algorithm 2 below. We note that while SOFTMAXBALANCE takes input sequence of key and value embeddings, it can nevertheless be used to compute the softmax normalization: simply run it the keys, with the corresponding value vector one-dimensional and all equal see line 11 BALANCEKV V. where COFTMAXBALANCE is called within the corresponding invocation of MERGEANDREDDUCE with value vectors as 1s. It's guarantees
to see
follows.
Theorem 3.3. Given sets K = (ki...., ki) C Rd, (ui,... ......., kn C R2, and failure probubility }n }. There exists randomized
> O, define to be the dataset of pairs C = [(k), (1),,)),- ..... (km, Un) algorithm, SOFTMAXBALANE
which outputs subset C" C, |C" < [C1/2, such that, for any

---

Algorithm 2 SOFTMAXBALANCEL Uz]s Meey Ivalue,
1: input: stream of I key-value embeddings (kz, Ug), radii rkey, rvalue: maxy ||kill2 S Tkey maxj |,yl|2 I Tvalues probability of failure 8.
2: R + exp(ricy/22 /2vd) Tvalue
3: f- 30 log(n/0)
4: Initialize zero vector + {0}
5: for from I and until the end of the stream do
6: y
(exp((ks,, (k)) vd) e3) kEDi
7: if Flyyn n|
c. R then FAIL
v"
8:
Pj +
20-12
+1 with probability Pj
9:
-1
o.w.
10: Add a new zero coordinate 7/++1 <
11: end for
if [[(K, 0.) : 7 = 1}1 S ((kg, v4) : 7 = 113 then
13:
output: {(ks, 84) 7)
14: else
15: output: [(ks, u.) 1h -13
16: end if
vector q Idd, with probability at least - 8,
(k:g)
(k. g)
Σ
exp
)
--- Σ exp
e
Vil
Vd
(AID)EO
(K.U)EC'
1gll3
|1ky112 2vd
(og(oonss/)) exp
exp
max jEini
mmx jE(n]
2vd
The runtime of SOFTMAXBALANCE is O((d + s)n?) and memory is ((dd + s)n).
The proof of the above theorem uses the breakthrough result [3] for the vector balancing problem, one of the main problems in discrepancy theory. Given a set of vectors ,i,..., kn the result of [3] produces subset of these vectors of at most half the size such that for any vector we have that Exec 9) ExetnJyc (k: q) with high probability. Our main contribution is algorithm for the vector balancing problem with respect to the function exp( (k: ///dd)o compared to (k, +) in the case of [3]. We defer the proof of Theorem 3.3 to Appendix A.2.
3.2 MERGEANDREDUCE.
As briefly mentioned above Section 3, MERGEANDREDUCE MAXBALANCE. The idea is to partition the stream of tokens into batches of size t. apply SOFT- MAXBALANCE to the batches to reduce the size of each batch by constant factor, and then repeat recursively see Fig. 2 in the appendix.
streaming version of SOFT-
we batch size & be about 1/s (see Theorem 3.4 below for the more precise setting), we obtain a streaming algorithm that approximates E/=1 exp( qj) Vdyu, at any point j the stream using total space Ö(dv oidddeer'/val the norms of key and query embeddings.
and runtime Ö ar* / /va/=") per step, where is upper bound on
As before, an important aspect is that MERGEANDREDUCE can handle value embeddings of dimen- sion not necessarily equal to that of key and query embeddings. Thus, when run on scalars W; 1 for all i, it can also be used to approximate softmax normalization at any point j in the stream. This is the main subroutine used in BALANCEKV V to approximate Aut(g, Kj, V,). Its pseudocode description is presented in Appendix A. and its proof is in Appendix A.3.2. The formal guarantees are
7

---

Relative Emor vs Compression Rate Llama
Relative Emor vs Compression Rate for Mistral
0.5
Layer BalanceKV
Layer BalanceKV
Layer Uniform
0.6
Layer Uniform
Relative Error
0.4
Layer
Relative Error
0.5
Layer BalanceKV
Layer 2 Uniform
Layer 2 Uniform
0.3
Layer 5 BalanceKV
0.4
5 BalanceKV
Layer 5 Uniform
Layer 5 Uniform
0.3
0.2
0.2
0.1
0.5
0.1
0.5
0.4
0.1
Compression Rate
Compression Rate
Figure I: Comparison of relative enrors across different layers of Llama-3 3.1-BBBBBBBB-B-- IB-Instuuct (left) and Ministral-Ohe Iww.t---nnttt----- 2410 (right) on TriviaQA dataset.
Theorem 3.4. For any r, > any of tokens (gi, ki, ui),.., (yns Kn; Un) where qj Kg Rd satisfy 11 [,kll2 S r, Ui E R2 for 8 S d suppose,
batch size = Oe
20 /Va /e) and compression rate 2 2-T with T = log(n/t).
Then MERGEANDREDUCE on input parameters t, r, r,,dd,a 8, outputs at every step j of the stream subsets of key-value embedding pairs C CO...., ....,CTTT. C C == ((k1,11),. v1),....... (Kn, bn)} such that, Zj = - 2t (k
exp
(k.qs)
dd
satisfies with probability at least 1 1/poly(n).
(ki 9n)
Σ iml
exp
Vd
Zj
Sej -
max |0ll2.
<E[n]
Total memory of the algorithm is sc 20 / /E), its j-th iteration runtime is Ódse 4r vr////////// /vd//2).
4 Experiments
this section we now present our experimental results. The full details of all sections as well as the experimental setup and implementation can be found in Appendix B.
4.1 Ablation Studies on Single Layer Attention Approximation
We evaluate the effectiveness of BALANCEKV for approximating attention in individual layers of [27] and I-sttrrutttt from LongBench [5]. Specifically, examine layers 1, 2, and and compare against independent uniform sampling key and value embeddings.
[52] on the TriviaQA dataset
Due to space limitations, provide the full experimental details in Appendix B. For each layer, we approximate attention for recent tokens using a compressed cache that retains fixed number of initial and recent embeddings, alongside intermediate ones selected via BALANCEKV V, and measure its relative eтror against exact attention. We vary the compression rate 2-T T E {1/2, 1/8, 1/16]. As shown in Fig. BALANCEKV consistently yields lower relative approximation error than approximating attention by uniform sampling past key value pairs across all settings, empirically validating its advantage as predicted by Theorem 3.1.
For fixed dataset and layer, also analyzed how the performance and runtime of BALANCEKV depend on the batch size and compression rate. More precisely, we repeat the single-layer attention approximation experiment TriviaQA and layers and 15 of Llama- 1aw.ea..------nnttueee for batch size (64, 128, 256) and compression rate 2-T E [1/2, 1 1/4, 1/8). The results are presented in Figure 3. As this experiment suggests, the quality of attention approximation increases as the size of the block doubles, while the runtime becomes slower also proven theoretically.

---

Method
yasper
multi
hotpotyaa
2wiki
gov
multinews
tree
unviaqu
Samssum
p.count
p.ret
loc
reporp
average
Qwen2.5.32B--
Instttu
Exact (Baseline)
44.56
50.65
69.
60.39
21.3
19.52
75.33
81.14
43.17
22.0
99.67
50.9
22
51.77
StreamingLLM
20.12
34.35
51.84
48.23
19.09
17.10
61.00
51.14
28.52
23.33
41.33
39.19
26.54
35.52
PyramidKy
46.33
67.78
$5.92
15.16
15.39
69.33
63.24
40.36
99.33
48.32
34.35
47.13
SnapKV
36.21
46.78
66.64
57.02
16.35
16.07
70.33
77.53
41.08
22.0
99.33
49.04
35.62
48.77
Uniform
39.28
43.82
57.84
23.10
19.50
73,00
81.63
39.70
22.00
92.00
44.97
32. 19
BALANCEKV
40.14
43.17
64.46
58.06
22.26
20.32
73.00
80.68
41.07
22.33
92.0
44.95
32.43
48.84
Qwen2,5-14B-Inet
(Baseline)
43.39
53.71
22.4
88.75
99.0
46.3
LL
20.73
32.62
49.93
42.39
21.63
18.69
59,67
74.96
29.57
11.67
63.0
46.16
32.25
38.71
PyramidKV
31.76
46.6
62.83
50.0
19.08
17.68
65.0
85.52
42.61
22.0
99.33
60.62
44.4
49.8
SnapKV
32.95
47.53
61.96
50.34
20.29
18.28
60.67
88.75
42.97
22.67
99.33
61.32
45.84
50.22
Uniform
37.07
41.69
61.11
49.96
29.18
22.67
71.67
87.89
40.34
22.0
84.33
58.0
42.57
49.88
BALANCEKV
37.02
41.96
61.74
50.9
29.26
22.64
71.67
88.1
41.14
23.67
87.67
58.64
43.63
50.62
Llama-3,
1888--
Exact (Baseline)
42.87
48.54
52.05
38.6
31.31
22.07
71.67
91.85
42.36
20.37
98.13
49.62
42.73
50.17
SIreamingI.IMM
20.65
30.71
39.14
32.43
23.10
18.70
58.00
83.87
28.85
20.36
97.26
33.69
30.46
39.79
PyramidKV
33.86
39.75
47. 12.12
35,96
20.03
17.78
63.67
90.76
40.21
20.4
98.97
45.25
39.51
45.64
SnapKV
33.91
42.55
49.09
36.13
20.48
17.67
62.0
91.7
40.23
20.33
98.86
46.7
39.86
46.12
Uniform
37.18
37.15
47.49
37.82
27.56
21.06
68.67
90.48
36.13
20.33
96.26
45.72
37.09
46.38
BALANCEKV
35.75
37.04
46.37
36.24
27.09
90.88
37.88
20.39
96.65
48.45
41.4
46.77
Table I: Comparison of various cache compression methods on LongBench-E using various models.
The best results among compression methods for each model are highlighted in bold.
4.2 End-to-end Evalution on LongBench
Next we evaluate BALANCEKV V on LongBench dataset [5], which tests long-context understanding across tasks like QA, summarization, few-shot learning, synthetic reasoning, and code completion. Specifically, we test a version of uniform length distribution (LongBench-F) During inference, we compress the key-value cache in the prefill stage using uniform compression rate of (approximately) 0.25 across all methods, while retaining all streamed embeddings during the decoding phase. We compare against StreamingLLMM [76], SnapKV [46]. PyramidKV [14], and uniform sampling (see Section 4.1), using their implementations from MInference [37]. The evaluation follows the Long- Bench protocol, using three pre-trained models at different scales including Llama-3. 1--8BB...1BBB--1 IB--nsttuc Qwen-2.55 14.5-14B 14B- Instruct and Qwen-2. 2.5--2..-32B 2.5-322--nnsteen Instruct. The results are reported in Table 1.
Notably, BALANCEKV consistently achieves the best overall performance among compression methods and across all models, demonstrating its effectiveness in preserving model quality for cache compression. The full experimental setup details can be found in Appendix B.2.
4.3 Needle-In-A-Hayee Benchmark
We evaluate BALANCEKV on the Srcedle-In---Haystee benchmark [39], comparing it against SnapKV, PyramidKV StreamingL.LM, and uniform sampling using Llama-3. 1-8B-BBBBB Instruct. The test challenges the model to retrieve specific sentence (the "needle"") embedded at an arbitrary position within a long context (the "haystack"). Following the setup in [29], we hide the needle at varying depths, from 0% to 100% of the total context length, across documents ranging from approximately 4K to 100K tokens. As in the previous experiments, all methods are evaluated under a fixed compression ratio of approximately 0.25.
To further enhance performance, we introduce an augmented version of BALANCEKV that deter- ministically preserves a small set of tokens whose key embeddings are strongly anti-correlated with the rest. The standard BALANCEKV procedure then applied the remaining tokens within each layer. As a result, BALANCEKV V achieves an average accuracy of 0.99, outperforming SnapK (0.83), Py (0.90), StreamingLL.M (0.31), and uniform sampling (0.90). Detailed heatmaps of performance across different context lengths and needle depths are in Figure 4 in Appendix B.3.
4.4 System Efficiency Metrics
We measure wall-clock times for both the prefill stage (including cache compression) and the decoding stage using a random input of length 16,384 tokens, followed by the generation of 1,024 tokens.

---

Results are averaged over 10 independent runs, with the minimum runtime reported to enhance robustness. The full results are provided in Table 2 in Appendix B.4.
All compression methods incur some prefill overhead compared to the uncompressed baseline (Fxact). While Streaming[.LM achieves the fastest decoding speed, it suffers from significantly lower accuracy (see Table 1). Among the remaining methods, BALANCEKV achieves the lowest prefill latency and consistently delivers the best trade-off between efficiency and accuracy. This demonstrates that our discrepancy-basse approach not only scales well in theory but also brings practical gains in end-to-end system performance, making a compelling choice for real-world deployment scenarios.
4.5 Additional Experiments
We additionally conduct the following experiments and provide their results in Appendix B.5 due to the space limitation.
1. Our main theorem (Theorem 3.4) relies upper bound of & norms of both query and key vectors. To validate this, we investigate the & norms of queries, keys, and values (QKV) on the TriviaQA dataset from LongBench [5] using -aama--..--B-IIntruuet Instruct. Specifically, analyze prompts in TriviaQA and compute the average & norms of all QKV vectors across all layers and attention heads during the prefill stage. The key findings are that all QKV norms consistently concentrate around some constants (15 for query, 15 for key, and 3 for value) with small confidence intervals (CT). Importantly, the norms remain stable across wide range of sequence lengths, suggesting that these norms do not grow with input sequence length.
2. We perform the evaluation of BALANCEKV and the uniform sampling when applied to the InternVL2.5-88 multimodal LLM for compression rates 1/4 and 1/16, for evaluation on the MS COCO image captioning dataset. The experiment was run on a NVIDIA A100 GPU with 80 GB VRAM.
3. We repeat the experiment in Section 4.2 in the extremely low compression rate regime on some of the datasets from LongBench [5]. More specifically, we compress the key-value cache in the prefill stage of inference using a uniform compression rate of (approximately) 0.8. 0.9, and 0.95 with uniform sampling as well as BALANCEKV. V, while retaining all streamed embeddings during the decoding phase. BALANCEKY demonstrates improved performance over uniform sampling across each of the compression rates and datasets.
4. We augment Section 4.2 by adding comparison to ClusterGen [83] using .lama-3.1-85 Instruct. The results are reported in Table 6.
Conclusion
We propose BALANCEKV, a token pruning method grounded in discrepancy theory. BALANCEKV enables approximate attention computation, which we both establish theoretically and validate empirically. To demonstrate the effectiveness BALANCEKV we conduct end-to-end experiments on a range of popular benchmarks and models of varying sizes. Finally, our work introduces theoretical problem of optimal streaming attention space complexity.
KV cache compression algorithm,
References
[I] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303308874, 2023.
[2] Joshua Ainslie, James Lee-Thorg Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages (895-49011 2023.
[3] Ryan Alweiss Yang P. Liu, and Mehtaab Sawhney. Discrepancy minimization via self- balaneing walk. Proceedings of the 53d ACM Symposium on the Theory of Computing (STOC 2021), 2021.
10

---

[4] Saleh Ashkboos. Amirkeivan Mohtashami, Maximilian L Croci, Bo Li, Pashmina Cameron, Martin Jaggi, Dan Alistarh, Torsten Hoefler, and James Hensman. Quarot: Outlier-free 4-bit inference in rotated llss. arXiv preprint Xiiv:24440004500 2024.
[5] Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong. Jie Tang. and Juanzi Li. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint 2ati::2308.14500 2023.
[6] Wojeiech Banaszezyk. Balancing vectors and gaussian measures of n-dimensional convex bodies. Random Structures and Algorithms 12 (1998), 351-360, 1998.
[7] Wojeiech Banaszezyk. On series of signed vectors and their rearrangements. Random Structures and Algorithms 40 (2012), 301-316, 2012.
[8] Nikhil Bansal. Constructive algorithms for discrepancy minimization. 51th Annual IEEE Symposium Foundations of Computer Science (FOCS '2010). arXiv: 2010.
[9] Nikhil Bansal, Haotian Jiang, Sahil Singla, and Makrand Sinha. Online vector balancing and geometric discrepancy. In Proceedings of the 52nd Anmual ACM Symposium on Theory of Computing (STOC '2020), 19Xiv: 1912.0330, 2019.
[10] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.055550 2020.
[11] Aline Bessa, Majid Daliri, Juliana Freire, Cameron Musco, Christopher Musco, Aécio Santos, and Haoxiang Zhang. Weighted minwise hashing beats linear sketching for inner product estimation. In Proceedings of the 42nd ACM SIGMOD-SIAACT-SIGR Symposium on Principles of Database Systems, PODS '23, page 169-181, New York, NY, USA, 2023, Association for Computing Machinery.
[12] Stéphane Boucheron, Gábor Lugosi, and Olivier Bousquet. Concentration inequalities. In Summer school on machine learning, pages 208-240. Springer, 2003.
[13] Vladimir Braverman, Avinatan Hassidim, Yossi Matias, Mariano Schain, Sandeep Silwal, and Samson Zhou. Adversarial robustness of streaming algorithms through importance sampling. Advances in Neural Information Processing Systems, 34:3544-3557, 2021.
[14] Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang. Junjie Hu, et al. Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling. arXiv preprint arXiv:2406. 02069, 2024.
[15] Moses Charikar, Michael Kapralov, and Erik Waingarten. A quasi-monte carlo data structure for smooth kernel evaluations. In Proceedings of the 35th ACM-SIAM Symposium on Discrete Algorithms (SODA '2024), arXiv:2401.00 025662, 2024.
[16] Moses Charikar. Similarity estimation techniques from rounding algorithms. In Proceedings of the thiry-fourth annual ACM symposium on Theory of computing, pages 380-388, 2002.
[17] Daniel Dadush, Aleksandar Nikolov, Kunal Talwar, and Nicole Tomczak- Tomezak-Jaegemman Jaggerman.. Balancing vectors in any norm. 59th Annual IEEE Symposium on Foundations of Computer Science (FOCS '2018), 2018.
[18] Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, et al. Deepseekmoe: Towards ultimate expert specialization in mnixture-of-expeer language models. arXiv preprint mxXiv:240.0660 2024.
[19] Majid Daliri, Juliana Freire, Christopher Musco, Aécio Santos, and Haoxiang Zhang. Sampling methods for inner product sketching. Proc. VLDB Endow. 2024.
[20] Majid Daliri, Juliana Freire, Christopher Musco, Aécio Santos, and Haoxiang Zhang, Sampling methods for inner product sketching, 2024.
[21] Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.0860) 086911 2023.

---

[22] Sanjoy Dasgupta and Anupam Gupta. An elementary proof of theorem of johnson and lindenstrauss. Random Structures & Algorithms, 22(1):60 65, 2003.
[23] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Gpt3, int8 (): 8-bit matrix multiplication for transformers scale. Advances in Neural Information Processing Systems, 35:30318 8-3032,, 2022.
[24] Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized llms. Advances in Neural Information Processing Systems, 36, 2024.
[25] Tim Dettmers, Ruslan Svirschevski, Vage Egiazarian, Denis Kuznedelev, Elias Frantar, Salch Ashkboos, Alexander Borzunov, Torsten Hoefler, and Dan Alistarh. Spqr: A sparse- sparse-quantized representation for near-lossless IIm weight compression. arXiv preprint urXiv:2306.03077 2023.
[26] Shichen Dong, Wen Cheng, Jiayu Qin, and Wei Wang. Qaq: Quality adaptive quantization for Ilm kv cache. arXiv preprint arXiv:2403.0400 2024.
[27] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle. Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Ilama 3 herd of models. arXiv preprint arXiv:2407.21783 2024.
[28] Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210.17325, 2022.
[29] Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannanch Hajishirzi, Yoon Kim, and Hao Peng. Data engineering for scaling language models to 128k context. arXiv preprint arXiv:2402. 10171, 2024.
[30] Yu Fu, Zefan Cai, Abedelkadir Asi, Wayne Xiong, Yue Dong, and Wen Xiao. Not all heads matter: A head-level kv cache compression method with integrated retrieval and reasoning. arXiv preprint arXiv:2410. 19258, 2024.
[31] Jianyang Gao and Cheng Long. Rabitq: Quantizing high-dimensional vectors with theoretical error bound for approximate nearest neighbor search. Proceedings of the ACM on Management of Data, 2(3):1-27, 2024.
[32] Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac"h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A frame- work for few-shot language model evaluation, 2023. https://gitae wtt.au...oooeeeeee weto//Eleutee www./////eeee//////// In-evaluatioon-arreee
[33] Mina Ghashami, Edo Liberty, Jeff M. Phillips, and David P Woodruff. Frequent directions: Simple and deterministic matrix sketching. SIAM J. Comput., 45(5): 1762-7792 2016.
[34] Insu Han, Rajesh Jarayam, Amin Karbasi, Vahab Mirrokni, David Woodruff, and Amir Zandieh. Hyperattention: Long-context attention in near-linear time. arXiv preprint 10.0022310.05869, 2023.
[35] Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length Ilm inference with kv cache quantization. arXiv preprint arXiv:2401.1800 18079999... 2024.
[36] Jianqiu Ji, Jianmin Li, Shuicheng Yan, Bo Zhang, and Qi Tian, Super-bit locality-sensitive hashing. Advances in neural information processing systems, 25, 2012.
[37] Huiqiang Jiang, YUCHENG LI, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin-Yew Lin, et al. Minference 1.0: Accelerating pre-filling for long-context IIms via dynamic sparse attention. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024.
12

---

[38] William B Johnson, Joram Lindenstrauss, and Gideon Schechtman. Extensions of lipschitz maps into banach spaces, Israel Journal of Mathematics, 54(2):129-338, 1986.
[39] Greg Kamradt. Needle in a haystack- ppressccre testing Ilms. Github Repository, page 28, 2023,
[40] Hao Kang, Qingru Zhang, Souvik Kundu, Geonhwa Jeong. Zaoxing Liu, Tushar Krishna, and Tuo Zhao. Gear: An efficient kv cache compression recipefor near-lossless generative inference of Ilm. arXiv preprint arXiv:2403.05527, 2024.
[41] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, arXiv preprint Stti:22001083666 2020.
[42] Junhyuck Kim, Jongho Park, Jaewoong Cho, and Dimitris Papailiopoulos.. Lexico: Ex- treme kv cache compression via sparse coding over universal dictionaries. arXiv preprint arXiv:2412.08890, 2024.
[43] Janardhan Kulkarni, Victor Reis, and Thomas Rothvoss. Optimal online discrepancy minimiza- tion. In Proceedings of the 56th Annual ACM Symposium on Theory of Computing (STOC "2024,, arXiv:2308. 01 4060,, 2023.
[44] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, pages 61 1-66,2 2023.
[45] Dacheng Li, Rulin Shao, Anze Xie, Ying Sheng, Lianmin Zheng, Joseph Gonzalez, Ion Stoica, Xuezhe Ma and Hao Zhang. How long can open-source Ilms truly promise on context length?, 2023. https: htapp:///huggigggg http///tggagg.....eee wwwp/nnffa.....oaanoae// co/wmmys////cccccc/ccccccssss h/tp///oggchhae oonghaat-b--v1....... 5-11..332..
[46] Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Llm knows what you are looking for before generation. arXiv preprini arXiv:2404. 14469, 2024.
[47] Lin, Jiaming Tang, Haotian Tang, Shang Yang, Xingyu Dang, and Song Han. Awq: Activation-aware weight quantization for I1m compression and acceleration. arXiv preprint arXiv:2306.00977 00978, 2023.
[48] Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of impor- tance hypothesis for 11m kv cache compression at test time. Advances in Neural Information Processing Systems, 36, 2024.
[49] Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. Kivi: A tuning-free asymmetric 2bit quantization for kv cache. arXiv preprint arXiv:2402.02750. 2024.
[50] Namiko Matsumoto and Arya Mazumdar. Binary iterative hard thresholding converges with optimal number of measurements for 1-bit compressed sensing. J. ACM, 71(5), October 2024.
[51] Jayadev Misra and David Gries. Finding repeated elements. Sci. Comput. Program., 2(2): 143- 152, 1982.
[52] Mistral AI team. Mistral ai, 2024.
[53] OpenAl. Introducing gpt-4o, 2024. http:://ppeaa.ccom//ddeeelllll-p
[54] Adam Paszke. Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Ginelshein, Luca Antiga, al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[55] Jeff M Phillips and Wai Ming Tai. Near-optimal coresets for kernel density estimates. Discrete and Computational Geometry, 63(4))867-887, 2020.
13

---

[56] Yaniv Plan, Roman Vershynin, and Elena Yudovina. High- -igeesiinnni estimation with geomet- ric constraints. Information and Inference: Journal of the IMA, 6(1):1-40, 2017.
[57] Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling transformer inference. Proceedings of Machine Learning and Systems, 5, 2023.
[58] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen, Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125, 2022.
[59] Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillierap. Jean- haptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint 20Xxv::20030055300 2024.
[60] Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman. Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 22500- 22500-22510, 2023.
[61] Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention- Fast and accurate attention with asynchrony and low-precision. arXiv preprint arXiv:2407.08608 2024.
[62] Noam Shazeer. Fast transformer decoding: One write-head all you need. arXiv preprint arXiv:1911.021500 2019.
[63] Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Ré, lon Stoica, and Ce Zhang. Flexgen: High-throughput generative inference large language models with a single gpu. In International Conference on Machine Learning, pages 31094-31116 316.. PMLR, 2023.
[64] Hanshi Sun, Wi-Wen Chang, Wenlei Bao, Size Zheng, Ningxin Zheng, Xin Liu, Harry Dong, Yuejie Chi, and Beidi Chen. Shadowky: Kv cache in shadows for high-throughput long-context Ilm inference. arXiv preprint arXiv:2410221465. 2024.
[65] Antropic Team. claude, 2024.
wwtp:/trrrrppp.ccoooooooeeeoe 3-
[66] FireFly Team. Adobe firefly, 2023.
[67] Llama3 Team. Llama3, 2024.
[68] Microsoft Copilot Team. Microsoft copilot, 2023.
https://ggthab..co//eaaarre
copilot.
[69] Midjoumey Team. Midjoumey, 2022. httt::://ww.mmijjorrrey.oo//aa/
[70] OpenAl Team. Sora: Creating video from text, 2024.
[71] Qwen Team. Qwen2.5: A party of foundation models, September 2024.
[72] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint urXiv:2307.09288 2023.
[73] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. 2017.
[74] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's trans- formers: State-of-the-art natural language processing. arxiv. arXiv preprint arXiv: 191000377, 2019.
14

---

[75] David Woodruff. Cs 15-859. Algorithms for big data lecture 11. https: wwwp://wwwggg/////////////aa/t// (al10/09--110////eaaa
//www. cs. cmu, edu/ http//fssss//////eeeee/eeeeeaass http://///..oooouu//////// Lecture_11s pot.outtttsoorrerhhcaa..... com, 2020.
[76] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint Stth::309.17455.. 2023.
[77] An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong, Haoran Wei, Huan Lin, Jialong Tang, Jialin Wang, Jian Yang, Jianhong Tu, Jianwei Zhang. Jianxin Ma. Jin Xu, Jingren Zhou, Jinze Bai, Jinzheng He, Junyang Lin, Kai Dang, Keming Lu, Keqin Chen, Kexin Yang. Mei Li, Mingfeng Xue, Na Ni, Pei Zhang, Peng Wang, Ru Peng, Rui Men, Ruize Gao, Runji Lin, Shijie Wang, Shuai Bai, Sinan Tan, Tianhang Zhu, Tianhao Li, Tianyu Liu, Wenbin Ge, Xiaodong Deng, Xiaohuan Zhou, Xingzhang Ren, Xinyu Zhang, Xipin Wei, Xuancheng Ren, Yang Fan, Yang Yao, Yichang Zhang. Yu Wan. Yunfei Chu. Yugiong Liu, Zeyu Cui, Zhenru Zhang, and Zhihao Fan. Qwen2 technical report. arXiv preprint arXiv:2407. 10671, 2024.
[78] June Yong Yang, Byeongwook Kim, Jeongin Bae, Beomseok Kwon, Gunho Park, Eunho Yang, Se Jung Kwon, and Dongsoo Lee. No token left behind: Reliable kv cache compression via importance-awan mixed precision quantization. arXiv preprint 241i::2402.889900 2024.
[79] Felix Xinnan X Yu, Ananda Theertha Suresh, Krzysztof M Choromanski, Daniel N Holtmann- Rice, and Sanjiv Kumar. Orthogonal random features. Advances neural information process- ing systems, 29, 2016.
[80] Yuxuan Yue, Zhihang Yuan, Haojie Duanmu, Sifan Zhou, Jianlong Wu, and Liqiang Nie. Wkvquant: Quantizing weight and key/value cache for large language models gains more. arXiv preprint arXiv:2402.12065, 2024.
[81] Amir Zandieh, Majid Daliri, and Insu Han. Qjl: 1-bit quantized jl transform for kv cache quantization with zero overhead. arXiv preprint MxXx:2400600488) 2024.
[82] Amir Zandich, Insu Han, Majid Daliri, and Amin Karbasi. KDEformer: Accelerating trans- formers via kernel density estimation. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th Inter- national Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 40605 40623. PMLR, 23-29 Jul 2023.
[83] Amir Zandieh, Insu Han, Vahab Mirrokni, and Amin Karbasi. Subgen: Token generation in sublinear time and memory. arXiv preprint arXiv:2402.06008, 2024.
[84] Tianyi Zhang, Jonah Yi, Zhaozhuo Xu, and Anshumali Shrivastava. Kv cache is 1 bit per channel: Efficient large language model inference with coupled quantization. arXiv preprint arXiv:2405.03917. 2024.
[85] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36, 2024.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Theorem 3.1 provides the main theorem for the performance of our algorithm.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contri- butions made in the paper and important assumptions and limitations. A No or NA answer this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long it is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Section 4.3 we point out that an enhancement of our main algorithm is needed for better performance on the Needle-In-a-Haasttcc experiments.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model nell-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faireess.
While the authors might fear that complete honesty about limitations might be used by reviewers grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play important role developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
16

---

Justification: Please refer to Section 3 and the proofs of the claims made there to find the full set of assumptions made and the full proofs.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Please refer to Section 4 and the links inside to find details of all experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, if the contribution necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to hosted model (e-g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation,
may be
While NeurIPS does not require releasing code, the conference does require all submissions provide reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g.. to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, described in supplemental material?
17

---

Answer:
[Yes]
Justification:
The code for all experiments is provided the supplementary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines pttttps://nips.cc//pllte
quides/CodeSubmi wwwilbmiissoooPoooiiooco for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
/
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparame- ters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Please refer to Section 4 for details regarding the experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification: The experiments conducted in Section 4 have been performed with error bars whenever applicable.
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar the standard deviation or the standard error the mean.
18

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report 2-sigma error bar than state that they have a 96% CI, if the hypothesis Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric eтror bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer:
[Yes]
Justification: Please refer Section 4 and the links therein find the details of all details.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't t make it into the paper).
9. Code of ethics
Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neur// hcpppbbcc//Ethhcccsiddeeeeeeee
Question:
Answer: [Yes]
Justification:
All studies conducted in this paper conform with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consideration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
The paper is fundamental in nature and has no direct path for societal risks or consequences. While the motivation comes from large deep learning models, which of course have myriad such potential ramifications, our paper does add to those ramifications.
Justification:
Guidelines:
The answer NA means that there is no societal impact of the work performed.
the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), faimesss considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
19

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mecha- nisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: As discussed in answer to the previous question, no specific safeguards are necessary for the work described in this paper.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: Existing assets used in adherence to their licenses and usage terms and with appropriate credit to their creators.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's creators.
20

---

13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
No new assets.
Guidelines:
The answer NA means that the paper does not release new assets. Researchers should communicate the details of the dataset/code/moodl as part of their sub- missions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL include anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, well as details about compensation (if any)?
Question:
Answer: [NA]
Justification:
No research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements your country institution) were obtained?
Answer: [NA]
Justification: No IRB approval needed for the research described in this paper.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage LLMs if an important, original, or non- standard component of the core methods in this research? Note that if the L.M is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
21

---

Answer: [NA]
Justification: We only used LLMs for visualizing results for submission and facilitating or running experiments.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, or non-standard components.
Please refer to our LLM policy htttps://neurips cc/Conferences/ /0t///022//2///// for what should or should not be described.
22

---

A Full Proofs
A.1 Proof of Theorem 3.1
Finally equipped with Theorem 3.3 and Theorem 3.4 we now state the proof of the main Theorem 3.1.
Proof of Theorem 3.1. Recall that BALANCEKV approximates attention by finding good estimations for the numerator and the denominator of attention separately. In line (8), erase those terms from the numerator, whose value vectors have sufficiently small lg norms. In what follows, we:
Bound the space and time requirements of BALANCEKV, as well as it's probability of failure. Each of the bounds readily follows from Theorem 3.4,
Bound the contribution of the erased terms to attention,
Show that BALANCEKV (or, more concretely, procedures MR-NUMERATOR,) approximate the rest of the terms of the numerator well,
Show that BALANCEKV (its subroutine MR- MR-DENOMINATOR) approximates the denominator of attention well.
We begin by analyzing the time/space requirements and probability of success of BALANCEKV. It never runs MRRRUMERATOR, for > (Xga("max) and i < log2(= n 11/2 Umax], so it never keeps more than log2( lo/2(vn/E) = O (n) of them. BALANCEKY at any step j performs one iteration of procedure - WERR-EEOMMNATORR one iteration of MR-NUMERATOR, with 2' 2 || 2 2'-1, computes the subsets Ki...., Kr, Vi...., VT and computes a function of the selected points in the subsets in line. Therefore, the runtime of BALANCEKV during one iteration is bounded by the maximum of the runtime of MERGEANDREDUCH during one iteration and time to compute the output. The latter is equal to its total memory. This maximum, by definition of L, is equal to O(Pe /va je²). The memory the algorithm is the union of memory of MR-NUMERATOR, for all i and MR- OR-DENOMINATOR, so the space complexity of the algorithm is Ö v de /e22"////e/ /E). The failure probability bounded by union bounding the failure probabilities of all instances of MERGEANDREDUCE and at most queries in the stream, and is equal to 1/poly(n).
Next, bound the contribution of the erased terms. Formally, let Umax( (J) = maxig ||ul||2 and define i(j)
max;
{
2
<
Vn Umax
j. Observe that
Σ
exp
(#2: Vid
Ui
Σ
KaS ||
exp
(R vd |/u/
||:112252101
by triangle inequality
E/m exp
(kr.g1) vdl
E exp
(kr.g1) vd
Σ Cis
exp
(krggh) vi
Il (2))))
by the definition of i(j)
2/7
E/=1 exp
(kr.g.) Vd
Umax
Kg. Kg. zj Vd
)1s
<
softmax
1IV,
(3) and
where the last inequality follows from the general inequality
|softmax
K Krgg Va
1/2
vn
||V, || 2 Umax-
We analyze the quality approximation of the denominator together with the quality of approxmmation of the numerator, both follow from Theorem 3.4. At time step j. procedure PIR-DENOMINATOO
23

---

returns subsets K1, "i....,Ky Kr such that
(Riss Vd
93)
T Σ 2
(ki Gz) vd
(ki Qj)
Σ exp AELA
2
S j. va E Σ exp
exp
vd
(K
(4)
as follows from Theorem 3.4 by plugging in scalars on =... 4 1. Next, define Pi.j = {k 2i : 1 S j. 2 ||oi|2 2 1. Intuitively, Pi aggregates all of the tokens proccessed by MR- MR-NUMERATOR; which appeated before time step j. MR-NUMERATOR; returns subsets V. ,...., V _ P., for all such that,
(k.og) Vd
T 2
(kgg) bi vd
Σ exp
Σ
Σ exp
U
(k,0)EPLs
1=0
(N.U)EV
S <P/Pual - 2
5 P
2
[K,9) qi
){{
exp
vî
vd
as follows from Theorem 3.4 observing that
|/1/| S 2'. The last inequality holds
because exp
(k.9) vd
2 ee-r/ve /v/ and, therefore, |epp
Kyng vi
2 vī v
Now, observe that |Vill S VE. Pi.j 2 and j Ei P.yll By the Sauchy-Schwand inequality.
Σ 1Pall-22
Σ PPsl 22:
Σ Pasl
ElPssh
By triangle inequality, the sum of the outputs of procedures MR- IR-NUMERATORR approximates the sum of the terms in the numerator of attention that were not erased in line (8) up to an additive error
exp
K2 Kg-0 di V
||
Σ
IP 1P s
exp
K Vd
||slle,
(6)
where the last inequality follows from Equation (5).
It remains show how the statement of the theorem follows from the derived bounds. Consider the following abstract derivation. Let u and u' be vectors such that || - u'112 S ar and let b and i/ be positive numbers such that
1 b
1
4
b.
Then, by application of triangle inequalities,
u
u'
<
||u - u' 112 + |uu|2
1
b
b
2 b'
b b
1
+ ||w||2
1
(7)
1 y
b
7 7
b
+ ||u||2
1 1-7
1- b.
From Equation (4),
1
1
1
$
E Eieti exp
(Rrne)) Va
>T 2 E ER exp
(kgg) vd
Eve(s) exp
vd
(8)
For simplicity of notation, let D C bi] denote the subset of indices of all tokens discarded
in line (8). Take ? = ΣiEby exp
(kr.gs)
D'
Σr 2 E Exk, exp
(k.az))
Vd
va
" EiEbIN exp (kk.g1) 1/. w' = Ei= 2 Σ Vi
E(R.oEE exp [(aaa) and, finally, a = 8
24

---

g exp
[V The first of the preconditions of our derivation holds by Equation (6). The
vd
2
second holds by Equation (8). Hence,
Σ T 2' Σ (AAW)EV exp
(R:g:)
Vd
EiebieD exp exp
(Niva,) vd
Vi
El 2 (KKDEE exp
[k.gg) vdd
EiEWs exp Ky Kg.8. gi vd
(keves? vi |1,
2
/sottmax
-IVVllF,
8
since ||u||2
D exp
( vd
V/s
S
exp
K
IVlle
vd
Finally, combining the above with Equation (3) via triangle inequality, we get
(K, 91 Vd
2
Σ
exp
1=0
(k.))V
2c -- E
Kg [gg vd
Attn(gj, K2 Vz))
S
softmax
1
T
(k, q1) Vd
Σ
Σ
exp
I=0
(K.))K!
By rescaling
c/4, get the desired approximation Equation (2).
A.2
Theoretical Guarantees of SOFTMANBALANCE
Algorithm Self-Balancing Walk introduced in [3] receives as input vectors u1..... Un and selects signs for them so that, for any direction, the signed sum of the vectors is balanced along that direction with high probability. The following theorem readily follows from theorem 1.1 in [3]:
Theorem A.1 (Theorem 1.1 in [3]). For any n, d E N, there exists a randomized algorithm which receives as imput a set of vectors U = {ui,.... un } The and a parameter 0. The algorithm
outputs (random) subset U' U such that, for any vector u E Rd, with probability at least
Σ NEU
(ui, u)
Σ T
(uis u)
s0 (0log(n/8)) max |uell2 1112-110 -luull))))) xE|m]
Algorithm 3 Self- Self- Balancing Walk ((,)),,r,))
1: input: stream of S n vectors uy, radius r: max, ||u 12 S r, probability of failure 6.
2: <- 30 30lggn//0)
3: U-, U+ +
4: for from and until the end of the stream do
5: if Eveu, (u, uz) Enev (u, ui) c. 2 then
6: Fail
7: end if
Pi 1/2 Eoco, (mas) -iuco (mmu))
8:
9:
Ei + + with probability ps, and e; +- - with probability P:
10: Us, + UUe U (kx
11: end for
12: if [U-| < |U_| then
13: output: U+
14: else
15: output: U_
16: end if
Proof of Theorem 3.3. Define for any I2d an embedding function (p(k):
(k/d0.2588
(kk))
Vi!
25

---

It is easy to see that for any two vectors k. q E R
(ki 9) vd
(p(K)), - exp
and for any & Ra
| vd
[lp(k) exp
Consider the set of vectors ((ki) & u,...., (p (kn) @ on Run the Belf-Balancing Walk algorithm the set of vectors p(ki ) 8 Ul,..., (kk,))))))))))))))) @ Un with failure parameter set to 6/s and denote by C" and C* C" the partition of C' returned by the algorithm. Observe that, even though vectors yp(s) @ us are infinite dimensional, Self-Balancing Walk still can be implemented. The algorithm never has to keep these vectors the memory because the only operation which requires the knowledge of the embeddings the inner product can be performed if we just store vector pairs {ks, b
(kigAig) ) Vd
(p(kk) & ui, (pky)) s vy) exp
(Us, Uz)-
Denote by 61, el....,6, C, the standard orthonormal basis in R By Theorem A.1, for any i [s] with probability - 5/8
Σ (s(k) 8 u, ((,)))0)))))))) 8e))
Σ (p(k) v. p(() 8ei)
(K.))EC
(9)
50 50 (og(nns//))
(ls(() so oll2 (1x(1) s ealle)
max (k.0)EC
and so with probability at least 1 all of the above inequalities hold simultaneously. To simplify [M
the right hand side, notice that [1)p(k) @ 2ll2 = exp
2v || and [|{(()) @ < exp
|
2va
Observe that for any (((k) u, (pg) @ e;) = (v²(),, ((()) [o] = exp
(k.q)
[o] where by
da
[v]c we denote the i-th coordinate of the vector Therefore, the left hand side of the expression above is simply the absolute value of the i-th coordinate of the vector
Σ exp
) V
(k,
Σ exp
)
(k, (1:0)) Vd
Vd
(k.
(R.P)EC
Thus, Equation (9) provides a uniform upper bound on the absolute values of coordinates of the above vector. Since the loc norm of a vector is the maximum of the absolute values of its coordinates,
(k:
(kigg) vd
Σ exp
Σ
exp
Vd
max (c[+]
Σ (p(()) @ (. 0c,)
Σ
(<((k) S (v.)) 0c)
(R.
||k1[2 2vd
|1
50 log((s/0))
max (K.
exp
||0||2
exp
2
26

---

Sofmaxiislance
N2 tokens
tokkets
takens
C
C
NZ tokens
<f tokeess
< Nokeens
o hosenss
< toiless
ttakens
Figure 2: Illustration of the tree structure of MERGEANDREDUCE
Finally, we go from bounding the (s norm a vector to bounding its l2 norm:
(kr
(k,
Σ
Σ
exp
exp
Vd
Vd
(A,D)EC"
(k.e)EC'
(k: 1)
(kg
$ V8
Σ exp
exp
vd
(A,O)EC"
00
||
11
0
V (o-_oo(s)/))
max (K.
exp
||0112
exp
2vd
2v7
A.3 MERGEANDREDUCE
A.3.1 Pseudocode for MERGEANDREDUCE
The pseudocode for MERGEANDREDUCH is presented
Algorithm 4 (MIRGEANDREDDCE((,,,.).....
1: input: stream of < n tokens (kj, 0j), batch size f, compression rate 2 precision parameter E.
2: Let SOFTMAXBALANCE be the algorithm as per Theorem 3.3.
3: Initialize i-th level subset C = 0, T, to empty
4: repeat
5: C" '" U (4kg.
6: if is not a multiple / then then
7: output C*) CO. ...., C
8: continue
9:
end if
/*Update subsets every steps*/
11:
p + j/t, i
12:
while p is an integer and until = T do
13:
C++1 <- Cr+1 U SOFTMAX Tkey, I'aluee 1/poly(1) (1))
14:
C 0-00
111
16:
p- p/2
17:
end while
18:
output C* CT CT
19:
until token stream ends
27

---

A.3.2 Theoretical Guarantees of MERGEANDREDUCEE
We now present the full proof of Theorem 3.4.
Theorem 3.4. For any T, O, any set of tokens k1, (1),),,................,,,,..,,) (gns (Onshin kng Un) where qj K2 Rdd satisfy |ai/l2: k S r, U/ I² for 8 S d suppose,
batch size t = ÖlVvss se / /v/// /2) and compression rate 2 with T = log(n/
Then MERGEANDREDUCE subsets of key-value embedding pairs C C ,",..., C C == vi),,.),.... (kn: vn)} such that,
input parameters t. r, d. s. outputs at every step j of the stream
2j = Ei (k, exp
(k:gg) v, satisfies with probability least 1 1/poly(n).
vd
Σ
j
exp
(vva))
(Kr.90)
144 2j
Sej rej.e-r9ve
max 11ll22- <E(n]
Vd
Total memory of the algorithm is Ö(dvse22* /va/e),
/va its j-th iteration runtime is Ö(se //// /v2
Proof. Let first consider the performance of the procedure at time steps which are multiples of t. Note that since in the statement of the theorem T = loga (n/t), condition until in line while redundant. Observe that any such j-th step the procedure is an online implementation of the following simple offline recursive algorithm on dataset {{ 01 (.,, uy}::
1. Set = j/t and = 1. Split the dataset (1k1, ei...., [ky, } into batches B B of size
2. While p is an integer:
Run SOFTMAXBALANCE on the batches B "i..,..., B
independently
If p is odd, store the output of SOFTMAXBALANCE on Bi- in C
For every I, merge the outputs of SOFTMAXBALANCE B and B into batch and store them in Bi,
Update - 1p/2], i +- i 1. Stop when p =
Therefore, will analyze space complexity and performance guarantees of the above offline algorithm.
Probability of success.
Note that our algorithm performs correctly if each of the calls to SOFTMAXBALANCE produces small error on each of the queries g (as in theorem Theorem 3.3). Throughout the stream, we make (nn/t) calls to SOFTMAXBALANCE, and we apply each to at most n queries, so, enough to require that that all SOFTMAXBALANCE have failure probability parameter d 1/poly(n).
Space complexity.
Observe that after each iteration step the number of batches decreases by a factor two. The maximum batch size is always bounded by t. This is because batch B1 at iteration of step 2 is union of SOFTMAXBALANCE B2-1) and SAATTAAAALANCEBBE
for B2-1 and Bu at iteration i i-1 1 of step and SOFTMAXBALANCE reduced the size of the dataset which it has been applied to at least by a factor of
The memory of the procedure is the collection of memory cells C and |C*| < I. Since at time step j at most logz()) < (n/z(n/t) memory cells are occupied, the total memory is bounded by (dt (ngt))))) = (dtT), which, using the Ô notation, is equal to dd 2 /Va /E).
Performance of the algorithm.
Define B = UB the data points which remained in the batches after iterations of step By triangle inequality.
28

---

i
7
(
5
(ki.g))
(ki,
2r
Σ
exp
Vdl
Σ
cxp
Vd
V/
(k)
2
E
exp
0- 2'
Σ
exp
Vil
vdl
(Ko))B++UUC
(k,d)EB
T-1
2'
exp
(ki (ki
Gi)
E
exp
(ki Uz)
vdl
(K+))B++UUC++)
(B:))ER((--UCC)))
We will refer the i-th summand (starting from on the right hand side as the error produced by the i + 1-st iteration of step 2. At the + 1-st iteration of step 2 we apply SOFTMAXBALANCE to p/22 batches Bi, Bi ... of size i, and we save the outputs of SOFTMAXBALANCE in batches Ch+1, Bi+1, B5+ Therefore, by Theorem 3.3 and triangle inequality, the error vector produced by the procedure at the + 1-st iteration of step 2 has la norm bounded by
O 0 2' Vo log(sn) Voggss
P 2
e" //a
max |15ll2 jE(n
0
(Vo (V8-log(on) log(sn) p ./V max lluslle)::
jEIn
since the error parameter 8 of instances of SOFTMAXBALANCE is set 1/poly(n). The norm of the total error of our procedure is bounded by
o (vos
V log(sn) T 2/va max jEin
||:]112]
By definition, = j/t. In order to ensure that the statement of the theorem is correct, the upper bound the l2 norm of the error vector of the procedure should be less than the desired error e-r//Vd max4e(n) ||u/||2:
o ( log(sn) T 2 "//aa
/va max |05/12 Sej e- j
max
|
t JE(n)
And, since by definition
22r2 V5..22"//0 V8 V5. (oog((m)) 0
1ogg
O
E
the above inequality holds.
Runtime during one time step. At worst, during j-th time step the algorithm has to launch SOFTMAXBALANCE (p) S log,(m// = times on batches of size t, so the runtime bounded by Td. In the notation, the runtime is equal to a /vā /E2).
As the final step, will analyze the performance of the procedure at time steps j' which are not multiples t. Define je [j/t]] t. Note that at any such time step the procedure simply saves the triplet (gr, ky, Uy) and outputs the sum the approximation 2j, such that
(ki, 1g) Vd
Σ exp
S --22/va max |o; 12. xc(m]
U,
2n
(1 Vd
and Σ exp
8/. From the above inequality,
29

---

(kis {3)
(kin
Σ exp
uu - Zjz +
Σ
exp
S
Eji e-"
mmxx |0 sEin
Vd
Vd
f=
as desired.
B Full Experimental Details
Experiments in Section 4.1 and Section 4.3 are performed on single NVIDIA A100 GPU with 80GB VRAM, and the rest on single NVIDIA RTX A6000 GPU with 48GB VRAM.
Implementation Detail. To enhance the practical performance of our algorithm. we implement BALANCEKV with parallel operations. Specifically, we consider the cache embeddings of length n and dimension b x (n/b) x d. Then, BALANCEKV V is applied in parallel to all blocks of length b. For cases where not divisible by b, we pad the embeddings with zeros. After sign assignment all embeddings in each block (i.e., line 9 in Algorithm 2), it is reshaped to its original length, and we strictly select n/2 embeddings, repeating this process for T iterations.
sequence of blocks with length and reshape them into a tensor of shape
B.1 Ablation Studies on Single Layer Attention Approximation
In this section we re-state with full details the single layer attention approximation experiments presented in Section 4.1.
We empirically evaluate the performance of BALANCEKV for approximating a single attention
layer, and compare it with independent uniform sampling. We use the pretrained Llama-3.1 ---88B Instruct [27] and Ministral-8B- -iw------nntttt-----
[52] and TriviaQA dataset from LongBench [5], and consider the 1st 2nd and 5th layers of the models for attention approximation.
For given a prompt with length n, we store the corresponding query, key. and value embeddings for all layers. Denote pair of embeddings in some layer by (q) kı, 01 vi)..... (gn, kin Un) and the goal is to approximate the attention Attn(gj- Kj. Vj) for the latest 256 queries, i.e. j E [n 256, n]. Specifically, we keep several first and recent tokens separately and apply BALANCEKV to the intermediate row vectors Kj. This is motivated by StreamingLLM [76] as important contexts are likely contained in the first and latest tokens. We retain the first 256 embeddings and the recent ones from n - 256 to j and our compressed cache contains tokens whose indices are in [256] USS U (n 256,. ... j where S C [257, 256] can be obtained from BALANCEKY V. We explore four compression parameters T E 21, 4} which reduces the cache memory by a factor of 2 2- . Let 2, be our approximation using BALANCEKV plus the recent and first few embeddings at the stream j E i - 256, n.]. We compute relative errors || At; Kj. Vi>le// Agtn(q, Ky. V2)lI for all j [n- 256, n.. batches, heads and input prompts in the dataset. We repeat this with 10 different random seeds and compute their average and standard deviations. We also compare our method to independent uniform sampling, in which replace the application of BALANCEKV V with sampling a 2-7 fraction of key and value embeddings with indices in [257, 256] uniformly at random. The results are reported in Figure 1.
Next we present the results of the ablation studies described in Section 4. I which demonstrate how batch size and compression rate affect the relative error in attention approximation for layers I and 15 for 2.1am-3.11-3 1--1--B-Inntru
B.2 End-to-End Evaluation on LongBench
We now provide the complete experimental details on the end-to-end evaluation in Section 4.2. We benchmark our algorithm on LongBench dataset [5], comprehensive collection of datasets designed to evaluate the long-context understanding capabilities of large language models. Specifically, we test a version uniform length distribution (LongBench-E), The benchmark consists of various long-text application scenarios, including single-document question-answering. multi-document summarization, few-shot learning, synthetic tasks and code completion. We use BALANCEKV to compress the key value cache generated in the prefill stage, and maintain all
30

---

Batch Size
1/2
1/4
1/8
1/2
1/4
1/8
256
0.0603
0.1190
0.1793
256
0.1036
0.1764
0.2655
128
0.0320
0.0624
0.0922
128
0.1082
0.1833
0.2741
64
0.0189
0.0349
0.0508
64
0.1137
0.1921
0.2858
(a) Layer Runtime (8)
(b) Layer Relative Error
1/2
1/4
1/8
1/2
1/4
1/8
256
0.3920
0.4505
0.5096
256
0.1107
0.1935
0.2798
128
0.3654
0.3951
0.4256
128
0.1121
0.1952
0.2813
64
0.3592
0.3753
0.3910
64
0.1141
0.1978
0.2845
(c) Layer 15 Runtime (s)
(d) Layer 15 Relative Error
Figure 3: Runtime and relative error for across different layers and block sizes. In each figure the rows are corresponding to various batch sizes and columns corresponding to various compression rates
streamed embeddings (dj: ky, v1)) during the token decoding/generraioon stage. This is because number of generated tokens is much smaller than the input sequence length. We set 256 and inputs.
2, achieving a consistent compression rate 0.25 across
We evaluate our method against several token-level key value cache compression schemes, including SLreamingLLLM [76]. SnapKV [46], and PyramidK [14] as well uniform sampling described in Section 4.1. We use their implementations from MInference [37], and configure their hyperparameters match a uniform compression rate with 0.25. We follow the same evaluation metries from [5]. We test them on Llama-3. 1-11---8ssssee I-8B-Instruct as well as bigger 14B and 32B parameter models Qwen-2.5- 14B-Instruct and 2reen-255-2B-Insttee
[77, 71J, with results summarized in Table
Our method consistently achieves the highest average performance among compression methods and across all models, demonstrating effectiveness preserving model quality for the cache compression. Notably, on the triviaqa dataset, it achieves near-exact scores compared to uncompressed baselines (c.g., 80.68 vs. 81.14 with Qwen2.5-32BB, highlighting its ability to retain high-quality information. We observe that uniform sampling performs competitively with our method and this result justifies that subset obtained from discrepancy theory has practical impacts on various LLM tasks.
B.3 Needle-In-A-Ie Heedle-InAAHaysstackk
In this section report the plots corresponding the Needle in a Haystack experiment described in Section 4.3. They are presented in Figure 4.
B.4 System Efficiency Metrics
In this section we present the prefill and decoding time numbers in Table 2 as described in the system efficiency experimental details in Section 4.4.
Method
Prefill Time (sec)
Decoding Time (sec)
Exact
3.032
37.769
SnapKV
3.755
40.426
PyramidKV
3.748
37,241
StreamingI.LM
3.681
40.276
BALANCEKV
3.662
38.054
Table 2: Minimum wall-clock runtime (in seconds) over 10 trials for prefill and decoding stages.
31

---

194
rnclsree Ama OEto anghtt/tatt tne reviece
Prose
Popus Lobemns
63
409
PAIA lang Liberin
113
MD0
Figure 4: Comparison of performance on Needle in Haystack task using Llama-3. 1.mma-3.1-BB-. 1--1--BB-Ta 1-8B-Instruct. The methods corresponding to figures from top bottom StreamingLLMM, SnapKV, PyramidKV, Unif. Sampling and BALANCEKV respectively.
B.5 Additional Experiments
1. In Table 3, we report the results of the experiment analyzing the l norms of query (Q), key (K) and value (V) embeddings described in Section 4.5. Due to space limit, we provide representative results in the below table from (randomly chosen) prompts of various sequence lengths. We note that the reported &a norms of keys shifted by their average, as opposed to the norms of the keys, because implementation of BALANCEKV shifts the keys by their
32

---

average before the compression. It is also easy to see that attention is invariant to the operation of shifting the keys by their average, so it is non-restrictive to assume that the average of the key values is zero.
Table 3: Statistics of norms of qkv embeddings for randomly chosen prompts from TriviaQA.
Query
Key Shifted
Value
Prompt
Seq Len
Mean
95% CI
Mean
95% CI
Mean
95% CI
10
2281
14.4512
0.0029
15.3451
0.0065
3.3398
0.0043
24
3131
14,6003
0.0025
15.6603
0.0058
3.3539
0.0036
30
3388
14.7406
0.0024
15.5152
0.0054
3.3419
0.0035
22
4230
14,8339
0.0021
15,6907
0.0049
3.3468
0.0032
5734
14.9259
0.0019
15.7149
0.0041
3.3482
0.0027
14
6616
14.9000
0.0017
15.6757
0.0038
3.3596
0.0025
4
6962
14.9743
0.0017
15.7346
0.0039
3.3450
0.0024
21
8041
14.9364
0.0016
15.7025
0.0035
3.3491
0.0023
26
17337
15.1437
0.0011
15.9531
0.0024
3.3654
0.0016
27
21274
15.2065
0.0010
15.8745
0.0022
3.3683
0.0014
2. In Table report the results of the multimodal task experiment described in Section 4.5.
Table 4: Comparison of BALANCEKV to uniform sampling on MS-COCO, using InternVL2.5-8B.
The bracket for every method contains the compression rate.
Method
Bleu 1 Bleu
Bleu
Bleu
METEOR
RougeL
CIDEr
Exact
0.795
0.629
0.476
0.351
0.291
0.580
1.255
BalanceKV (1/4)
0.794
0.628
0.475
0.351
0.290
0.579
1,251
Unif (1/4)
0.794
0.629
0,476
0.350
0.290
0.578
1.247
BalanceKV (1/16)
0.789
0.622
0.468
0.343
0.286
0.573
1,221
Unif (1/16)
0.789
0.619
0.465
0.340
0.284
0.571
1.207
3. In Table 5, we present the results of the end-to-end experiment LongBench in the extremely low error regime, described in Section 4.5. We note, that etremely low error regime corresponds to the low compression rate regime, and therefore performance of BALANCEKV V in the low compression rate regime. For compression rates of 0.8, 0.9 and 0.95, randomly select dataset from LongBench and apply both uniform sampling and BALANCEK V to achieve the desired compression rate. More specifically, if we wish to compress KV cache to 1 of it's original size, select its subset of size 2a, compress it by factor of 2 with either uniform sampling or BALANCEK V and keep the rest exactly.
experiment is equivalent to exploring the
Table 5: Comparison of BALANCEKV uniform sampling in LongBench the extremely low-error regime.
Compression Rate
Dataset
BalanceKV
Uniform
Baseline
0,8
HotpotQA
50.2
48.4
51.9
0.8
TriviaQA
91.6
86.3
91.6
0.9
MultiFielQQA
47.5
44.9
47.8
0.9
Qasper
42.3
39.6
43.1
0.95
LCC
49.3
45.7
49.5
0.95
P.Count
20.7
20.1
20.7
4. In Table 6, we present the results of the end-to-end evaluation of ClusterGen [83] as the experimental setup in Appendix B.2 and compare its performance to both BALANCEK and exact attention.
33

---

Method
qusper
multi
hotpoiiua 2wiki
multinews
Iree
triviaqa
samusum
p.ccunt
p.ret
Ice
repoop
average
Llama-3.1-88
Exact (Baseline)
42.87
48.54
52.05
38.6
31.31
22.07
71.67
91.85
42.36
20.37
98.13
49.62
42.73
50.17
ClusterGen
33.93
42.31
50.85
37.24
21.16
19.31
67.67
90.82
39.49
20.20
96.57
47.23
39.26
46.62
BALANCEKV
35.75
37.04
46.37
36,24
27.09
20.84
69.0
90.88
37.88
20.39
96.65
48.45
41.4
46.77
Table Comparison of ClusterGen, BalanceKV and exact attention on LongBench- using L1 ama- L-8B-Instruee The best results among compression methods for each model are highlighted in bold.
C Lower Bound
In this section, we prove the lower bound on the space complexity of an algorithm approximating the Attn(:, K, V) function. More formally,
Theorem 3.2. Suppose that 22 < d. Any streaming algorithm which input (1k1,21)... {kns h. l1q|22, ||kall2 S r. outputs 24 satisfying Equation (2) with probability 0.999 has space complexity Ω (minf ft, (app((////)))
The proof will be reduction to the well-known INDEX problem.
C.1 Reduction the INDEX Problem
Definition C.1 (The INDEX problem). Alice gets a bit string x ~ Unif(O, 13" and Bob gets i Unif[n]. Then, the goal is to compute fzz, z; on Bob's end with single message m from Alice. Denote by Rg R the public coin one-way communication complexity of computing a function f(z y) with error probability at most 6: Alice holds x, Bob holds y, they share source of random bits and Alice sends a single message to Bob, after which he must output the correct answer with probability at least 6.
Theorem C.2 (Proven in [75]).
B 2 2(n).
2/3
Proof of Theorem 3.2. Let e be the small constant such that 122/3 IRpub (INDEX) 2 c- n.
Assume the contrary to the statement of Theorem 3.2 that there exists a streaming algorithm of space complexity const. min (eexpx222///// /Vd) for any sufficiently small constant const. We d will show that given a string of length d exp((222 /d)) Alice can solve the INDEX problem as follows. She instantiates such an algorithm with const < <c/C for a sufficiently large constant C, gives it as input carefully selected set of keys and values, and sends the state of its memory to Bob. Bob, on his end, can determine whether any randomly drawn bit ~ Unif[n] equals or with probability 0.8 by issuing corresponding (carefully crafted) query to the streaming algorithm and observing its output. Thus, if the streaming algorithm uses small space, we get a contradiction with Theorem C.2, and therefore obtain proof of Theorem 3.2.
The reduction. Suppose Alice's ' input to the INDEX problem is a bit string æ E {0, 13" of length n min L. dcxp(222 /V)) Using public coins, Alice and Bob jointly generate n/d key vectors
ki,...,k,,/ Unif Va' Va function ग [n]] + [n/ d] x [d] which randomly partitions the n bits into groups of size d, and n random signs oI...., 0. Unif(-1,
Let (i)) E [n/d] be the first component of (() and (()2 E [d] the second component of π (1). Let e Ed be the standard orthonormal basis in Rdd We build the dataset of key-value pairs in the following way. We associate with the i-th bit the key vector k; = kx())) and the value vector 2 o, (x)()))
Define
U (ik, v.) xx, 1}
34

---

the set of key-value pairs corresponding to entries in Alice's bit string Alice instan- tiates the streaming algorithm for approximating Attn(,, K, V) with space complexity min 22 d expppr2////1eeeee V with a big enough constant C which we specify later and sends the state it's memory before reading q to Bob. When Bob receives the message, he uses it to approximate Attn(q:, K, V) where q: kg If the value written in the only non-zero coordinate of vs is larger than exp(r"/va)
10 max((exp((=/va]1101> Bob reports that the i-th bit of Alice's string is equal to 1, and otherwise
Analysis of the reduction.
Proof sketch. Before moving to formal proofs we briefly outline the main idea of the analysis. Observe that, by the choice of key vectors, any exp( /://vd)) for j * is in expectation insignificantly small compared to exp((ks, /d)/vd) sometimes we will even refer to these terms as
"noise". This statement is formalized in Lemma C.4. Therefore, when Bob computes an approximation to Attm (q: K, V), he will observe a large value in the coordinate where U; is non-zero if (ki, oi) E U and small value otherwise.
Lemma C.3. E
Unu - Sa jal } (exp((C(z,
((1) for any constant C.
Proof.
E
- (exp(C(ax,))/VV)
Unir -V.. -2
= 1
2 expp((r> /23/2 exp( exp(-cc>/d812)) cosh
(Cr² d
d3/2
(Cr2)
S CCppC2///4/ S cosh
d3/2 I exxp(C²r+/dF) < exp(C²)
where we used the assumption that r2/d /
Lemma C.4. Fix i [n), select qi ki Let [U& be the number of key-value pairs in U \ (ki vi) whose value vector has non-zero entry in the same coordinate as V;.
If (ki, vi E U then with probability >
1000
U
Σ
exp
(ki.es) vd
(o, vz)
2 (xpp)//v)-0O((U]]]
- -O(IT))
Otherwise, with probability
U.
1000
Σ (A
exp
(k, 9s)
(v, vz)
s so <v(VV))
Vd
Proof. We prove both statements using Chebyshev's incquality.
In the first case, i.e. when (ks, v.) U, the sum contains the term exp( (ks, 9:)/ Vd) cxp(r2/v4 /Vd). and otherwise it does not. It therefore remains to upper bound the absolute value of the sum
(kr.9t) Vd
(k, qi) Vd
X = Σ
Tk exp
(v. (p))))
Σ
Th exp
(o,us). us),
(R
EEU (R
(R..)
0% ~ Unif(-1, 1}. which effectively introduces "noise" in Bob's estimate of whether 2; We upper bound this sum now
35

---

B[X] 0, and Var(X) = |U:] Var
<lexp( v\/vd)) - O ((U))) because
a.y-Uniff va.
Var s.y~Uuitf- Ja v
(exp( (z. v\ vd)) s E
u.ur~Unff ( =
(exp((((xx))) v\()) s exp(4),
by Lemma C.3. We therefore get by Chebyshev's incquality
Pr |x| 1000 |
1
1000
|U]
Therefore, with probability
1
1000
(ki (K) Vd
Σ cxp
(v, 2/)) 2 exp(r>/Va) 1100
[R
In the second case, the entire sum equals Exev, on exp( q)) /V). As shown above,
Pr | x 1000 1000
|
S
T000 To
Hence, with probability 1 1000
[0
(k, qk) qs) S <1000 U
Σ (K.Dee
(v, vz)
exp
Vd
Corollary C.5. Suppose bits iia form a group that is, m(ii) = T(12) = _.. m(aa)
Then all Ui....., Vta have different non-zero coordinates, and therefore Σ S
Therefore, by the union bound argument, the conclusion of Lemma C.4 holds for all bits which form one group simultaneously with probability 0.999.
Lemma C.6. Fix a bit With probability 0.98 the following holds:
1. The error of the approximating algorithm in the only non-zero coordinate of 02 is bounded by
e
||ootmax -2)122 (0Vll))VVllr) (VIlr)
2. If the i-th bit is I then
(ke, qx) Vd
V vd
Σ exp
(u) 2 exp(r2/Va)
HR.UtEU
and
(k.9.)
V
Σ (K.U)EU
exp
(u, 0;)
50
Vd
Vd
otherwise.
Proof. We may think that the process of generating the dataset and the approximating streaming algorithm has the following order: first Alice and Bob jointly generate the partition æ, the key vectors generate Bob's input
ki...., kn /a and the value vectors Ui....,U, all using public randomness. position E [x] we generate pair Unif(n/ d], ~ Unif[a] and declare = π - (a, b). We may assume that is chosen before the datasets K and V are generated, and b after.
36

---

Before b is drawn, the key vector K of = (a, is already defined, as well as the datasets K, V and U. Alice can therefore already apply the streaming algorithm and Bob can already apply to 4i kg Therefore, the error vector which the streaming algorithm yields when applied to di k; also defined before is known.
Clearly, there are no more than 0.0001 coordinates in which the error of approximation exceeds 10000- - ||softmax q)ll [[V p. Since every value vector has only one non-zero entry, there va
more than 0.0001 d coordinates where at least 10000 of value vectors from U have non-zero value. We call all coordinates which are neither these two groups safe. From the above, at least 99% of the coordinates choosing the coordinate in which v; is non-zero. Therefore, with probability 0.99 over the choice of b the only non-zero coordinate of 0; is safe.
safe. Recall that b ~ Unif[d], and choosing is equivalent to
At the same time, similarly to Lemma C.4, by Chebyshev inequality, if U3 C U the set of all key- value pairs in U whose value vector has the same non-zero coordinate Uy then with probability 0.999 the i-th bit is then
((hsa)) Vd
) (u, vz) 2 exp(r>/Va) (VViil)
Σ exp (K,u)E
and
Σ
exp
(1.
(e, vz)
s s
vd
otherwise.
By union bounding over these two events, we get that the statement of the lemma is correct with high constant probability.
Conclusion of the proof. Let T C U be the of all pairs from U with the same key as (ki, v:)- Since Alice's string is drawn from Unif[-1,17 with probability 0.9 | 2 0.4 d. This because every bit in the same group k belongs U with probability 1/2.
Observe that by Chebyshev inequality, with high probability 0.999, the denominator of softmax (K ..) lies in range
|19 (u'x/exp(+//Vay 1 exp/(r/) 20 17]] 200111]
This is because every summand in the denominator, except for exp((k,, q.) Vd), is dis- tributed as exp((e,))/Vd)) a,y Unif Va vaj and the expectation and the vari- ance of this distribution, as shown in Lemma C.3, ( This range contained in (max(de"" [U})), (ma ( We will denote the denominator as D.
Similarly, by Chebyshev inequality, with probability 0.999 the numerator of softmax ( q:) lies in
[ -x4--exp(222/V + 1U1 expp22x2(2222/// + 200 10|
which, since |V| S (2x (202/2222//)) /V)) is bounded by 2000 /Vd
Suppose that the i-th bit is 1. Then,
37

---

When 1 2 Vder"/va, by selecting [0] de //d for some enough constant C' the /vv
value written in the only non-zero coordinate of o; least e=²iva eva
and at most 100 D
D
otherwise, as follows from Lemma C.4; 100 D
When < vde" "dde"/va /vá by selecting [U] = 2² for some big enough constant C' the value written in the only non-zero coordinate of U; is at least 2//aa E-D-Va otherwise, as follows from Lemma C.4.
S.D.Va' and most
The error which the approximator can have in the non-zero coordinate of is bounded by
10000 ||softmax(K 2)1l2 [| s 10000 & /20 Vde Iva
Vd
Vd
D
as shown in Lemma C.6. Below, we show that this error is smaller than the gap between .//v
40 m (va and the value written in the coordinate, which means that, even though the approximator introduces some error, Bob is still capable to tell whether the i-th bit is 1 or 0.
When 2
vde /va, the gap between the value written in the coordinate and A/vl
o 3iva maxi /v,,1101 is at least 1006 max de / vvv0 and the error
10000s
/20- /20-vde"* lvva
lva
vdeep(2r2 /vdd) S 1
er\/va
V
vd
D
2000 max (de /vd [0, 2000 max de /ve.
by an appropriate choice of C.
When Vde /va the gap between the value written in the coordinate and 2/ is at least 1000 max and the error
3/v max ( 1v,,111
10000c
/20 Vdee Iva
exp(r2/ve 2/v// VUUS 2000 max(der"/V2 //11 2000 max(de [UB
exp(12/ /vd)
Vd by appropriate choice of C.
38


Reviews:
[{"review_id": "dU0LQVgVMH", "reviewer": "Reviewer_BB75", "summary": "This 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.", "questions": "N/A", "limitations": "N/A", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### 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.", "was_revised": true, "final_justification": "All concerns have been resolved.", "rebuttal": "Thank you for carefully studying our submission and the additional, very useful suggestions and comments!\n\n1. 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\n\nResponse: 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. \n| Compression Rate |    Dataset     | BalanceKV | Uniform | Baseline |\n|------------------|----------------|-----------|---------|----------|\n| 0.8              | Hotpotqa       | 50.2      | 48.4    | 51.9     |\n| 0.8              | Triviaqa       | 91.6      | 86.3    | 91.6     |\n| 0.9              | Multifieldqa   | 47.5      | 44.9    | 47.8     |\n| 0.9              | Qasper         | 42.3      | 39.6    | 43.1     |\n| 0.95             | Lcc            | 49.3      | 45.7    | 49.5     |\n| 0.95             | P.Count        | 20.7      | 20.1    | 20.7     |\n\nExtremely long context regime is also a very interesting direction, and we will include it in future work.\n\n\n2. 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\n\nResponse: 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.\nIn 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. \nWe 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.\n\n| Prompt ID | Seq Len | Q Mean  | Q 95% CI | K Shifted Mean  | K 95% CI | V Mean  | V 95% CI |\n|-----------|---------|---------|----------|---------|----------|---------|----------|\n| 10        | 2281    | 14.4512 | 0.0029   | 15.3451 | 0.0065   | 3.3398  | 0.0043   |\n| 24        | 3131    | 14.6003 | 0.0025   | 15.6603 | 0.0058   | 3.3539  | 0.0036   |\n| 30        | 3388    | 14.7406 | 0.0024   | 15.5152 | 0.0054   | 3.3419  | 0.0035   |\n| 22        | 4230    | 14.8339 | 0.0021   | 15.6907 | 0.0049   | 3.3468  | 0.0032   |\n| 5         | 5734    | 14.9259 | 0.0019   | 15.7149 | 0.0041   | 3.3482  | 0.0027   |\n| 14        | 6616    | 14.9000 | 0.0017   | 15.6757 | 0.0038   | 3.3596  | 0.0025   |\n| 4         | 6962    | 14.9743 | 0.0017   | 15.7346 | 0.0039   | 3.3450  | 0.0024   |\n| 21        | 8041    | 14.9364 | 0.0016   | 15.7025 | 0.0035   | 3.3491  | 0.0023   |\n| 26        | 17337   | 15.1437 | 0.0011   | 15.9531 | 0.0024   | 3.3654  | 0.0016   |\n| 27        | 21274   | 15.2065 | 0.0010   | 15.8745 | 0.0022   | 3.3683  | 0.0014   |\n\n\nThe 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."}, {"review_id": "TtKLs2jyHj", "reviewer": "Reviewer_pBkK", "summary": "The 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.", "questions": "1. 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?", "limitations": "there is no \"limitations\" section.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "based on the strengths and weakness I previously wrote, other reviewers' comments, and  the authors response. I would like to keep  my positive original rating.", "rebuttal": "Thank you for carefully studying our submission and the additional, very useful suggestions and comments!\n\n--> 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.\n\n\nResponse: 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. \nIn 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.\n\n-->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\nResponse: 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. \n\n\nQuestions:\n1. 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?\n\nResponse: 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).  \nStudying 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. \n\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\n\nResponse: 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.\n| Method              | Bleu_1 | Bleu_2 | Bleu_3 | Bleu_4 | METEOR | RougeL | CIDEr |\n|---------------------|--------|--------|--------|--------|--------|--------|--------|\n| Exact               | 0.795  | 0.629  | 0.476  | 0.351  | 0.291  | 0.580  | 1.255  |\n| BalanceKV (¼)       | 0.794  | 0.628  | 0.475  | 0.351  | 0.290  | 0.579  | 1.251  |\n| Unif (¼)            | 0.794  | 0.629  | 0.476  | 0.350  | 0.290  | 0.578  | 1.247  |\n| BalanceKV (1/16)    | 0.789  | 0.622  | 0.468  | 0.343  | 0.286  | 0.573  | 1.221  |\n| Unif (1/16)         | 0.789  | 0.619  | 0.465  | 0.340  | 0.284  | 0.571  | 1.207  |\n\nFor 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. \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\n\nResponse: 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."}, {"review_id": "fKcfe46SGx", "reviewer": "Reviewer_w4SJ", "summary": "The 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).", "questions": "1. 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.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: \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?", "was_revised": true, "final_justification": "This paper provides an interesting way to save memory usage in transformers. The theory results are interesting and further verified by some experiments, although their theoretical techniques have been extensively used in theoretical computer science (for different purposes). I recommend a borderline accept for these reasons.", "rebuttal": "Thank you for carefully studying our submission and the additional, very useful suggestions and comments!\n\n-->What is a good estimation of r in practice? A large r will weaken the results.\n\nResponse: 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. \n\n-->This paper gives error upper bound for a single attention layer, but the error could compound when we have more heads and layers?\n\nResponse: 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. \n\nQuestions:\n\n1. 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.\n\nResponse: 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.\n\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.\n\nResponse: 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. \nMergeAndReduce 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$.\n\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\nResponse: 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."}, {"review_id": "xiX7NFRa86", "reviewer": "Reviewer_VYzU", "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.", "questions": "Please see weaknesses", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The strengths of this paper can be summarized as follows:\n\nFirst, 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.\n\nSecond, 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.\n\nThe weaknesses of this paper can be summarized as follows:\n\nFirst, 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? \n\nSecond, 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.", "was_revised": true, "final_justification": "I thank the authors for providing these clarifications, which have addressed all of my concerns. Since my initial rating was already high, I choose to maintain my original score.", "rebuttal": "Thank you for carefully studying our submission and the additional, very useful suggestions and comments!\n\n1. 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?\n\nResponse: Thanks for the question, please refer to point 2 of the response to Reviewer BB75 where this is discussed.\n\n\n2. 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.\n\nResponse: 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."}, {"review_id": "of5wcFvE07", "reviewer": "Reviewer_CpPu", "summary": "This 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.", "questions": "1. 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.", "limitations": "Yes, though I believe that a big more extensive discussion on limitations and open questions would be a good addition to the paper.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### 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).", "was_revised": true, "final_justification": "The paper contains a very nice algorithm for KV cache compression based on discrepancy minimization. The theoretical contributions are very nice and solid, and the experiments performed are insightful and promising. Few small revisions and clarifications in the main text are required to improve clarity, as per the discussions.", "rebuttal": "Thank you for carefully studying our submission and the additional, very useful suggestions and comments!\n\n-->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.\n\nResponse: 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.\n\n-->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?\n\nResponse: 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.\n\n-->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.\n\nResponse: 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.\n\nWe 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. \n\n|               Method | qasper  | multifieldqa_en | gov_port | multnews | trec   | triviaqa |\n|----------------------|---------|------------------|----------|-----------|--------|----------|\n| Exact                | 42.87  | 48.54           | 31.31   | 22.07    | 71.67 | 91.85   |\n| Subgen             | 33.93  | 42.31           | 21.16   | 19.31    | 67.67 | 90.82   |\n| BalanceKV       | 35.75  | 37.04           | 27.09   | 20.84    | 69.00 | 90.88   |\n\nAs a result, our BalanceKV outperforms Subgen on all datasets but “multifieldqa_en”. We will add complete experiments in our final draft.\n\n-->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.\n\n\nResponse: 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. \nRegarding 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.\n\nQuestions:\n\n1. I have not heard of the term \"streaming complexity\" before? I suppose it refers to space complexity?\n\nResponse: 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’’. \n\n2. 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?\n\nResponse: 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. \n\n3. What if the number of generated tokens is not known in advance? How would you set T?\n\nResponse: In practice, the maximum number of generated tokens is a property of the model and, therefore, it is known in advance.\n\n4. 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?\n\nResponse: 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.\n 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.\n\n5. 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?\n\nResponse: 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})$. \nUltimately, 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.\n\n6. What is the \"prefill\" stage? It is not clear what this stage is architecturally in the experiments you are mentioning.\n\nResponse: ''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''. \nIn 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."}]

Author rebuttal:

Meta-review: Although computing attention in transformer-based language models requires $O(nd)$ space to store the keys and values, for context length $n$ and embedding dimension $d$, this paper shows that these bounds can be bypassed when the goal is only to approximate attention, by using tools from discrepancy theory and the streaming model. The algorithm is further complemented with a lower bound and experiments. Initial concerns about the assumption that the $L_2$ norms of keys, values, and queries are bounded were well-addressed by the authors' response.

To strengthen the potential impact of the work, additional action items for the authors include:
- Clarifying the MR-Denominator subroutine, which is currently not described in sufficient detail for independent understanding
- Incorporating the new experiments described in the author rebuttal phase
Decision: Accept (spotlight)