Attention under Bounded Key-Query Similarity: Space Complexity and Transformer Anisotropy

Introduction

The high memory consumption of modern large language models (LLMs) is a major challenge in their deployment, as the decoding stage of transformers requires storage of all tokens seen so far. The growing body of work on Key-Value (KV) cache compression has been proposed to address this challenge. These algorithms rely on innovative ways of approximating the attention function of the context using small space, either by quantizing the vectors, or evicting some of them to retain the most informative ones. In both cases, the goal is to approximate the value of the attention function:

where $K=(k_1,\ldots,k_n)\subset\mathbb{R}^{d}$ is the set of all keys, $V=(v_1,\ldots,v_n)\subset\mathbb{R}^{d}$ is the set of corresponding values and query $q\in\mathbb{R}^{d}$ is the query.

We consider randomized algorithms for token eviction that are query-oblivious. That is, before seeing the specific query, the algorithm constructs a weighted summary by selecting a random subset of token indices $I\subseteq[n]$ together with positive weights $w_i$ (and evicts the other tokens). Given a query $q$, the algorithm returns

For every fixed query $q$, we require

The probability is over the randomness used to construct the summary. The query is fixed independently of this randomness; we do not require simultaneous success for all queries. Assume that $\lVert k_i\rVert\le r_k$ for every $i$ and $\lVert q\rVert\le r_q$. We will assume $\lVert v_i\rVert\le1$.

Query-oblivious weighted token eviction (a.k.a. coresets) has been a subject of multiple recent theoretical investigations, leading to (almost) tight bounds on the worst case space complexity Haris & Onak [9]; Liberty et al. [13]; Chen et al. [4], see Table 2 for details. These results show that an $e^{r_k r_q/\sqrt{d}}$ dependence in the space bound is sufficient and necessary. While these results are close to optimal in the worst case, some of their implications significantly deviate from what one sees in practical instances:

Reference	Model / regime	Upper bound	Lower bound
Kochetkova et al. [11]	Coreset size (tokens)	$\widetilde{O}\left(\sqrt{d}e^{2R^{2}}\right)$	$\widetilde{\Omega}(1)$
Haris & Onak [9]1	Space complexity (bits)	$\widetilde{O}(nd)$	$\Omega(nd)$
Liberty et al. [13]2	Coreset size (tokens)	$\widetilde{O}\left(\sqrt{d}e^{R^{2}+\log R+o(\log\log R)}\right)$	$\widetilde{\Omega}\left(\sqrt{d}e^{R^{2}}\right)$
Chen et al. [4]3	Space complexity (bits)	$\widetilde{O}\left(e^{(1+o(1))R^{2}}\right)$	$\widetilde{\Omega}\left(e^{(1-o(1))R^{2}}\right)$
This work4, $b-a\le\frac{1+b}{2}$	Coreset size (tokens)	$\widetilde{O}\left(e^{(b-a)R^{2}}\right)$	$\widetilde{\Omega}\left(e^{(b-a-o(1))R^{2}}\right)$
This work4, $b-a>\frac{1+b}{2}$	Coreset size (tokens)	$\widetilde{O}\left(e^{\frac{1+b}{2}R^{2}}\right)$	$\widetilde{\Omega}\left(e^{(\frac{1+b}{2}-o(1))R^{2}}\right)$

1. The space lower bound of $e^{r_k r_q/\sqrt{d}}$ can be much larger than the size of the input itself, given that the values of $r_k,r_q$ are often on the order of $15-20$ (see, for example, Table 3 in Kochetkova et al. [11] and our Figure 9) and $d=128$; in such cases, the upper bound becomes vacuous.

2. The hard distribution of the keys used to prove the lower bounds is isotropic, whereas in practice, keys and queries are known to be anisotropic, concentrating in a relatively narrow cone of the high-dimensional space Godey et al. [8].

3. The worst case hard instances have the queries align perfectly (i.e., have a high dot product) with one of the keys; in contrast, in practice, the key-query cosine similarity is small in absolute value once the attention sink is excluded and stored separately (see e.g., Mazaré et al. [16] and our Figure 1).

Vector (PDF)

Vector (PDF)

Our results

Reflecting this discrepancy, we study the space complexity of attention approximation under the following assumption:

$(\star)$ Every key/query cosine similarity lies in some range $[a,b]$, where the magnitudes $\lvert a\rvert,\lvert b\rvert\ll1$, i.e., no query can align significantly with any key.

Before we state our main theorem, we need to introduce some notation. Fix radius parameters $r_k,r_q$. Suppose $\lVert k_i\rVert\le r_k$ for every $i$, and fix a pair of parameters $a,b\in[-1,1]$. We take the admissible query set to be

Let $\mathrm{OPT}_{\mathrm{Attn}}$ denote the optimal worst-case number of retained tokens subject to

for every fixed admissible $q$. The algorithm knows keys, values, and parameters $r_k,r_q,a,b$, but not the particular query in $Q$. All norms are Euclidean, and all logarithms are natural. We use $\widetilde{O}$ to suppress factors polylogarithmic in $n$.

Our main result shows that importance sampling is a nearly optimal token eviction policy for attention approximation. Surprisingly, for a wide range of parameters (see Figure 2) all importance weights can be set to $1$, i.e. the most basic uniform sampling is optimal. Formally, we define two regions in the space of parameters $(a,b)$, see also Figure 2 for an illustration:

Our main result is as follows.

Theorem 2 is achieved by uniform sampling when the key/query cosine similarities are small ($a < 0 < b$ and $b-a \le\frac{1+b}{2}$), and by an importance sampling scheme for larger cosine similarities. Interestingly, even though uniform sampling is by no means the best token eviction policy, it is quite competitive with some earlier works on token eviction for KV cache compression, especially query-independent ones. See, for example, Table 1 in Kochetkova et al. [11] for a comparison of SnapKV and BalanceKV with uniform sampling. The recent work of Wang et al. [18] provides more evidence for this phenomenon. Figure 2 helps explain this from the perspective of theory: the empirical $(a,b)$ values for widely used models lie in or close to the uniform sampling regime.

Model	Layer 4	Layer 8	Layer 16	Layer 24
Llama 3.1 8B	$-0.584 \pm0.148$	$-0.725 \pm0.098$	$-0.681 \pm0.076$	$-0.688 \pm0.105$
Qwen3 8B	$0.068 \pm0.037$	$-0.146 \pm0.034$	$0.024 \pm0.044$	$-0.143 \pm0.080$
Ministral 3 8B	$-0.724 \pm0.047$	$-0.628 \pm0.086$	$-0.613 \pm0.078$	$-0.658 \pm0.045$

In addition to the theorem statement, the proof of the theorem provides insights into another empirical phenomenon. Specifically, the anisotropy of the key/query space that has been empirically observed in the literature Godey et al. [8] appears naturally in our lower bound instance. In our hard input distribution on key and query vectors, the keys and queries are similarly distributed except that the keys are positively shifted along a specific direction $e_0$ while the queries are negatively shifted along $e_0$. This construction, which is critical to establishing a matching space lower bound, is indeed anisotropic: the keys have a large positive projection on a common direction, and the queries have a large negative projection on the same direction, similarly to distributions observed in practice (Table 1, Figure 5). See § 2.2 for more details.

Vector (PDF)

Related work

Related work on worst-case space complexity of attention approximation. Prior work (Haris & Onak, [9]; Liberty et al., [13]; Chen et al., [4]) has established nearly tight asymptotic bounds on the space complexity of approximating attention in the worst case. The algorithmic landscape is diverse, with approaches drawing on sketching, random sampling, discrepancy theory, and kernel-density estimation. Provable attention approximation has also led to implemented KV cache compression methods. SubGen (Zandieh et al., [20]) combines online key clustering with $\ell_2$ sampling of values to obtain approximation guarantees that exploit clustering structure, and evaluates its implementation on long-context question-answering tasks. BalanceKV Kochetkova et al. [11] develops a streaming compression algorithm based on geometric discrepancy and vector balancing, with both theoretical guarantees and empirical evaluations of attention approximation and end-to-end generation on long-context benchmarks. A comparison is given in Table 2. We note that the worst case setting is crucially different from our setup in that queries can be arbitrarily well aligned with key.

Related work on KV cache compression. We mention the lines of the most relevant work.

KV subset selection. A prominent approach to KV-cache compression retains a subset of the original key-value pairs. H$_2$O combines accumulated-attention heavy hitters with recent tokens Zhang et al. [22], while Scissorhands uses token importance that persists across decoding steps Liu et al. [14]. StreamingLLM keeps initial attention sinks and a sliding window of recent tokens Xiao et al. [19]. SnapKV selects KV positions per head before generation using a window at the prompt’s end Li et al. [12]. PyramidKV splits a fixed cache budget across layers Cai et al. [2], while Ada-KV splits it across heads using an attention-output error bound Feng et al. [7]. These approaches exploit historical attention, prompt observations, or layer/head structure, rather than establishing optimal retained-token complexity under an explicit restriction on future key-query similarities.

Compressed KV representations. Another line of work reduces storage per cached vector. KIVI quantizes keys per channel and values per token Liu et al. [15]. KVQuant combines distribution-aware quantization, pre-RoPE key quantization, and separate outlier handling Hooper et al. [10]. TurboQuant uses random rotations for data-oblivious vector quantization, with guarantees on mean-squared and inner-product distortion Zandieh et al. [21]. Palu caches low-dimensional intermediate representations from low-rank factorizations of key/value projections Chang et al. [3]. Cartridges learns context representations through self-study Eyuboglu et al. [6]. These representations fall outside our weighted-subset model, so our retained-token lower bounds do not apply to arbitrary quantized or low-rank encodings.

Random KV eviction. Wang et al. [18] propose a method that reserves cache space for the input prompt and the most recent tokens, then fills the remaining slots by uniform random sampling, independently across KV heads at each eviction event. Their method is competitive with importance-based policies on reasoning benchmarks, providing further evidence that random sampling can be effective in practice.

Uniform sampling regime

Upper Bound

The upper bound in the small cosine similarities regime is shown by a simple approach: sampling tokens uniformly.

Consider the small cosine similarities regime of Theorem 1, where $-1 \le a < 0 < b \le1$ and $b-a \le\frac{1+b}{2}$. Our assumptions are $\lVert k_i\rVert\le r_k$, $\lVert q\rVert\le r_q$, and

We write $R^2 = r_k r_q/\sqrt{d}$. Set $m = \left\lceil\frac{96e^{(b-a)R^2}}{\varepsilon^2} \right\rceil$. If $m \ge n$, retain the full dataset. Otherwise, sample $m$ tokens independently and uniformly with replacement. Writing $J_1,\ldots,J_m$ for their

Vector (PDF)

indices, return the estimate

The intuition behind uniform sampling is that no individual token can account for too large a fraction of the attention mass. Indeed, for every $i$ and every admissible query $q$,

Thus, the attention denominator is at least $ne^{aR^{2}}$, whereas the contribution of any individual token is at most $e^{bR^{2}}$. Consequently, each token accounts for at most an $e^{(b-a)R^{2}}/n$ fraction of the attention mass.

We defer the full proof of Lemma 1 to Section B. We use Chebyshev’s inequality to bound the attention denominator, and use Markov’s inequality to bound the numerator of the error.

Lower Bound

Recall that in our model (2), the token-eviction algorithm (randomly) selects tokens and assigns them nonnegative weights before seeing the query. We will prove lower bounds against such algorithms.

Vector (PDF)

To see this, choose $u_1,\ldots,u_M$ independently and uniformly from $\{-D^{-1/2},D^{-1/2}\}^{D}$, so that each vector has unit norm. Hoeffding’s inequality and a union bound show that the probability of any pair satisfying $\lvert\langle u_i,u_j\rangle\rvert> \xi$ is at most $M^2e^{-D\xi^2/2} < 1$ under the assumed bound on $D$. We give the full proof in Section C.

In order to keep all key-query cosine similarities in $[a,b]$, queries and keys are anisotropically shifted away from each other along a fixed direction $e_0$. A pair of keys and queries in the same group $k_i,q_i$ are (a weighted version of) $u_i \pm e_0$ where the sign on $e_0$ differs for the key and query (along with a normalization term to set their respective norms to $r_k,r_q$).

Each query $q_i$ has similarity $b$ with its own group’s key $k_i$, while the near orthogonality caps its similarity close to $a$ with the other groups’ keys $k_j$. The exponential in attention amplifies this gap, giving the query’s own group almost all the attention mass. We also choose nearly orthogonal value vectors, so a convex combination of other groups’ value vectors cannot well-approximate attention. The algorithm must therefore retain a key-value pair from each group with constant probability, otherwise it would fail with probability more than $1/3$ on some group’s query.

The resulting construction holds an intriguing connection to patterns in real data. In Figure 5, we see that queries and keys are separated along a common direction and otherwise hold a somewhat similar distribution though with different scales for queries and keys. These qualities also appear in our hard instance (Figure 4).

Lower bound construction

Goal. We construct keys and queries whose scaled inner products are $bR^2$ for matching pairs $(k_i,q_i)$ and $a(1-o(1))R^2$ for nonmatching pairs $(k_j,q_i)$, where $R^2=r_kr_q/\sqrt{d}$. In Section C.1, we use this gap to make each query concentrate its attention on copies of its matching key, forcing an accurate summary to retain many key groups.

Construction. Let $e_0,e_K,e_Q$ be the first three standard basis vectors. Choose $m$ unit vectors satisfying Lemma 2 in dimension $D=d-3$, with pairwise inner products bounded in magnitude by $\xi=\sqrt{\frac{6\log(2n)}{D}}$. Embed them into the last $d-3$ coordinates of $\mathbb{R}^d$, and denote the resulting vectors by $u_1,\ldots,u_m$. Define

where $\alpha\approx b-a$ and $\beta\approx-a$. Specifically, we set

We have $\alpha\ge0$ since $b\ge a$. Since $a<0$ and $b$ are fixed and $\xi=o(1)$, for sufficiently large $n$ we have $\xi b<-a$, so $\beta>0$. Moreover, $\alpha+\beta=2\alpha-b\le b-2a\le1$. Therefore, $1-\alpha-\beta\ge0$ and the keys and queries are well defined with norms exactly equal to $\lVert k_i\rVert=r_k$ and $\lVert q_i\rVert=r_q$. The components along $u_i$ and $e_0$ serve the purpose of achieving the desired inner products while the padding along $e_K$ and $e_Q$ increases the norms of the keys/queries to fix the denominator of cosine similarity.

Verifying the inner products. For every $i,j\in[m]$, $\frac{\langle k_j,q_i\rangle}{\sqrt d}=(\alpha\langle u_j,u_i\rangle-\beta)R^2$. For matching pairs, $\alpha-\beta=b$ gives

For nonmatching pairs, $\lvert\langle u_j,u_i\rangle\rvert\le\xi$ gives

Constructing the value vectors. Apply Lemma 2 a second time to obtain $m$ unit value vectors $v_1,\ldots,v_m\in\mathbb{R}^d$ satisfying

Since $m\le n$, the assumption $d=\omega(\log n)$ ensures that this is possible. The final proof of the coreset size lower bound using these key, query and value vectors is given in Section C.

Importance sampling regime

Upper Bound

When the interval $[a,b]$ is wide, the exponential in attention can make a few tokens much more influential than the rest. To account for this, uniform sampling would need to retain many tokens. We thus use importance sampling instead. Unlike uniform sampling, which treats all tokens equally, importance sampling gives higher sampling probabilities to tokens that can contribute more to attention. This is a standard idea in randomized linear algebra through leverage-score sampling Drineas et al. [5] and in coreset construction through sensitivity sampling Braverman et al. [1].

Since we sample before seeing the query, we want to sample a token with probability proportional to the largest possible attention it could receive. This quantity is its sensitivity: for a nonempty query set $Q$ containing all admissible queries, define

Unfortunately, these sensitivities may be difficult to compute exactly. Instead, it suffices to use sampling scores $s_i\ge\sigma_i$, i.e., upper bounds of each sensitivities. In the sampling procedure, we sample each token $i$ with probability $\pi_i=s_i/S$, where $S=\sum_{i=1}^{n}s_i$. We set the number of samples proportional to $S$, so smaller sampling scores yield a smaller sample size. This leaves two tasks: compute valid sampling scores $s_i \ge\sigma_i$ and show that their sum $S$ is small.

We formalize the sampling procedure in Algorithm 1. For each sampled token $i$, we retain the corresponding key-value pair $(k_i,v_i)$ together with a positive weight. If we sample the same token more than once, we keep one copy and add the weights together.

Algorithm 1 Attention compression by score-based sampling and reweighting

Require: key-value pairs $(k_1,v_1),\ldots,(k_n,v_n)$, accuracy $0 < \varepsilon < 1$, and sampling scores $s_1,\ldots,s_n > 0$

1: $S \leftarrow \sum_{i=1}^{n} s_i,\quad m \leftarrow \lceil 96S/\varepsilon^2\rceil$
2: if $m \ge n$ then
3:   return $\{(k_i,v_i,1):i\in[n]\}$
4: end if
5: $\pi_i \leftarrow s_i/S$ for every $i\in[n]$
6: Subsample $m$ tokens using probabilities $\pi_i$ ▷ with replacement
7: $c_i \leftarrow$ number of times token $i$ was sampled
8: return $\{(k_i,v_i,c_i/(m\pi_i)):c_i>0\}$ ▷ one copy per token

We refer to the set returned by Algorithm 1 as the summary. Given a query $q$, the summary produces an estimator

where $I:=\{i:c_i>0\}$ denotes the set of retained token indices. The denominator is positive, since the summary is nonempty and all retained weights are positive.

We first state the approximation guarantee of Algorithm 1. Using a number of tokens proportional to $S$, the estimator from Equation (13) approximates the attention with high probability.

We defer the full proof of Lemma 3 to Section D. Similar to the proof of Lemma 1, we use Chebyshev’s inequality to bound the attention denominator, and use Markov’s inequality to bound the numerator of the error.

Lemma 3 reduces the problem to constructing sampling scores $s_i$ whose sum $S$ is small. To this end, we first bound the total sensitivity $\sum_{i=1}^{n}\sigma_i$ of any query set for which all key-query inner products lie in a fixed interval. This condition is satisfied by our query set $Q$. However, $Q$ need not be convex, making it difficult to compute sampling scores. We therefore partition the keys and possible query norms into small ranges and, for each pair of ranges, replace the cosine-similarity constraints by fixed inner-product bounds. This partitions the problem into a collection of smaller convex relaxations on which we can approximate the sensitivities. Applying the total-sensitivity bound to these relaxations allows us to control the sum $S$. These two steps are formalized in Lemmas 4 and 5.

At a high level, the two terms in Equation (14) come from two different ways of controlling the total sensitivity. The first term $h-\ell$ follows from the fact that $\sigma_i \le\frac{e^h}{ne^\ell}$. For the second term, we pair each key $k_i$ with a query $q_i$ attaining its sensitivity and assign weights to these pairs proportionally to their sensitivities. Writing $\bar{k},\bar{q}$ for the weighted mean key and query, using Jensen’s inequality, the proof gives

Averaging these two bounds yields $\log\sum_i \sigma_i \le(R^2+h)/2$. The full proof is deferred to Section D.

It remains to construct the sampling scores.

The full proof is deferred to Section D.

We now combine the approximation guarantee (Lemma 3) with the sampling score bound (Lemmas 4 and 5) to prove the upper bound in the importance sampling regime.

Proof. By Lemma 5, the sampling scores satisfy $s_i \ge\sigma_i$ and $S \le128e^2\kappa^{-2}e^{\kappa R^2}$, where $\kappa=\min\{b-a,(1+b)/2\}$. Since $a,b$ are fixed, $\kappa^{-2}$ is a constant. Substituting into Lemma 3 and absorbing $\kappa^{-2}$ into the implicit constant proves the upper bound.

Lower bound

The lower bound in the importance sampling regime uses the same construction and analysis as in the uniform sampling regime (§ 2.2). The only change in the theorem statement is that $m$ must satisfy

The full statement (Theorem 4) and proof are given in Section E.

Open Problems

The following questions naturally arise from our work:

• Can we provably explain why transformers learn this geometry? In particular, under what assumptions on the data and training dynamics can we prove that small key–query cosine similarities and negative alignment between mean keys and queries emerge?

• Do real KV caches have other structure that our model misses? Our hard instances share some geometric features with trained models, but real caches may have additional properties that make them easier to compress?

ACKNOWLEDGMENTS

JC,YF, PI, were supported in part by the NSF TRIPODS program (award DMS-2022448) and Wistron Corporation. Ekaterina Kochetkova and Boris Prokhorov were supported by the Swiss National Science Foundation grant 10004935 - Sublinear Algorithms for Non-linear Structures (Project MINT)

AI USE STATEMENT

In this work, we used generative AI tools to provide critical ingredients for proving mathematical claims in the importance sampling regime, assist in the writing of proofs, propose or refine hypotheses and implement methods.

We have not used generative AI tools to help formulate mathematical claims, develop theoretical models or conceptual frameworks, design or provide feedback on research methodology or experiments, interpret results, support qualitative and thematic data analysis. Generating synthetic data sets, cleaning and reformatting datasets and assisting with translation are not applicable to this work.

Additionally, we used generative AI tools for creating scientific figures, suggesting experimental parameters, creating software code, drafting parts of the paper, sourcing/searching for information, editing to improve readability, proposing a title for the paper.

We have reviewed all AI-assisted work. We checked LLM-generated parts of the proof and rewrote it ourselves. We take responsibility for the final content of this work, including text, claims or artifacts produced with the aid of generative AI.

Preliminaries

Standard Inequalities and Identities

We record several standard inequalities and identities used in the proofs.

Convex Optimization

We record the convex optimization facts used to compute the sampling scores. The optimization guarantee below is stated in the oracle model.

The convexity statement follows from convexity of log-sum-exp and composition with an affine map; see Vandenberghe & Boyd [17]. The gradient formula follows by differentiation.

Proof. The function $x\mapsto\lVert Ax+c\rVert-\langle u,x\rangle-\beta$ is continuous and convex. Thus its non-positive sublevel set is closed and convex, and these properties are preserved by intersection. Compactness follows from closedness and boundedness.

Proof. Since $C$ is compact and $f$ is continuous, $f$ attains its minimum on $C$. An exact minimizer is, in particular, an additive-$\eta$ solution.

Proofs from sec:upper-uniform

for uniform samples $J_1,\ldots,J_m$. We write the attention error vector as

We will use Chebyshev’s inequality to bound the denominator, and use Markov’s inequality to bound the norm of the numerator. By Equation (5), we have $D\ge e^{bR^2}$. Moreover, since $J_1,\ldots,J_m$ are independent,

where $J$ is uniform on $[n]$. Thus, by Chebyshev’s inequality (Lemma 7),

On the other hand, we can bound the expected squared norm of the numerator as

Hence, by Markov’s inequality (Lemma 6), the probability that the numerator has norm larger than $\varepsilon D/2$ is at most $4e^{bR^2}/(m\varepsilon^2D)\le1/24$.

Therefore, with probability at least $1-1/96-1/24>2/3$, we have simultaneously $\widehat{D}\ge D/2$ and numerator norm at most $\varepsilon D/2$, which implies $\|\widehat{\operatorname{Attn}}(K,V,q)-\operatorname{Attn}(K,V,q)\|\le\varepsilon$.

Proofs from Section 2.2

Here we prove a standard lemma which is used in the lower bound about packing nearly orthogonal vectors in high dimensions.

Proof. We will prove the result by a standard application of the probabilistic method. For all $i\in[M]$, choose each coordinate of each vector $u_i$ independently and uniformly from $\{-D^{-1/2},D^{-1/2}\}$. Then, $u_i$ is a unit vector. Note that for any $i\ne j$, $\langle u_i,u_j\rangle$ is the sum of $D$ independent random variables with mean zero and bounded in magnitude by $1/D$. Hoeffding’s inequality (Lemma 8) gives $\Pr[|\langle u_i,u_j\rangle|>\xi]\le2e^{-D\xi^2/2}$. A union bound over the $\binom{M}{2}$ pairs upper bounds the failure probability by

Therefore, such a set of vectors must exist.

Proof of thm:lb-attn

Here we present the proof of the lower bound in the uniform sampling regime. Refer to the main body in § 2.2 for the construction of key, query, and value vectors.

Proof. Each group will contain $t=\lfloor n/m\rfloor\geq1$ copies of the same key-value pair, with one additional copy in the first $s$ groups, so that $n=tm+s$ and $0\leq s<m$. For $i\in[m]$, let $N_i=t+1$ if $i\leq s$ and $N_i=t$ if $i>s$. The attention instance $K,V$ will have $N_i$ copies of the key-value pair $(k_i,v_i)$. By the setting of $t,s,N_i$, there will be exactly $n$ key-value pairs. As $t\geq1$, $N_j/N_i\leq2$ for all $i,j\in[m]$.

At distinct query $q_i$, let $\alpha_i$ be the attention mass on the copies of $(k_i,v_i)$. Bounding $\langle k_j,q_i\rangle-\langle k_i,q_i\rangle$ by Equation (9) and Equation (10), we get

Recall the condition $m\leq e^{(1-2\xi)(b-a)R^2}/32$. Then, $1-\alpha_i\leq2me^{-(1-2\xi)(b-a)R^2}\leq\frac{1}{16}$.

The attention output assigns mass $\alpha_i$ to $v_i$ and total mass $1-\alpha_i$ to the other value vectors. Since all value vectors have unit norm, the triangle inequality gives

Let $H_i$ be the event that the coreset retains at least one copy of $(k_i,v_i)$. On $H_i^c$, the nonnegative coreset weights imply that

Since $|\langle v_i,v_j\rangle|\leq1/8$ by Equation (11) for every $j\neq i$, we obtain

Since $\lVert v_i\rVert=1$, the Cauchy-Schwarz inequality and Equation (15) give

On $H_i^c$, we have already shown that $\langle v_i,\widehat{\operatorname{Attn}}(K,V,q_i)\rangle\leq1/8$. Since $v_i$ is a unit vector, another application of the Cauchy-Schwarz inequality gives

Hence success at the fixed query $q_i$ implies $H_i$. Since answering $q_i$ accurately requires retaining a copy of $(k_i,v_i)$, the success guarantee Equation (3) implies $\Pr[H_i]\geq2/3$ for every $i\in[m]$. Since each retained pair represents at most one group,

By definition, $\operatorname{OPT}_{\mathrm{Attn}}$ is the optimal worst-case number of tokens retained by an algorithm, so, from above, $\operatorname{OPT}_{\mathrm{Attn}}\geq2m/3$.

Proofs from sec:upperimportance

Lemma 3 (Approximation guarantee). Suppose $\lVert v_i\rVert\leq1$ and $s_i\geq\sigma_i$ for every $i\in[n]$. Then Algorithm 1 returns a weighted attention summary of size at most

For every fixed query $q \in Q$, its estimate satisfies

where the probability is over the sampling in Algorithm 1.

Proof. If the algorithm retains all tokens, the estimate is exact. Otherwise, fix $q \in Q$ and write

Thus $\operatorname{Attn}(K,V,q)=Y/Z$. The sampled denominator and numerator in (13) are

Since each index $i$ is sampled with probability $\pi_i$, we have $\mathbb{E}\widehat{Z}=Z$ and $\mathbb{E}\widehat{Y}=Y$. The sensitivity upper bound gives

Using $\pi_i=s_i/S$,

Similarly, using independence and applying Lemma 12 to a centered sample, we obtain

where we used $\|v_i\|\le1$ and (16).

Using Chebyshev’s inequality (Lemma 7) for $\widehat{Z}$ and Markov’s inequality (Lemma 6) for $\|\widehat{Y}-Y\|^2$, followed by a union bound, we obtain

hold together with probability at least $1-\frac{32S}{m\varepsilon^2}\ge\frac{2}{3}$, since $m=\lceil96S/\varepsilon^2\rceil$.

On this event, $\widehat{Z}\ge(1-\varepsilon/4)Z$. Also, $\|Y\|\le Z$ because $\|v_i\|\le1$. Therefore,

where the last inequality uses $0<\varepsilon<1$. The algorithm retains at most $\min\{n,m\}$ tokens, which gives the stated size bound. $\square$

Proof. In this proof, we replace each key $k_i$ by $\sqrt{r_q/r_k}\,k_i$ and each query $q$ by $\sqrt{r_k/r_q}\,q$. This preserves all inner products and hence all sensitivities, while giving both keys and queries the common norm bound $\sqrt{r_k r_q}=R d^{1/4}$. We use the same notation for the rescaled keys and query set.

Since $\mathcal{Q}$ is compact and the attention weights are continuous, for each $i \in[n]$ there exists a query $q_i \in\mathcal{Q}$ attaining the sensitivity of token $i$. We let $q_i$ be a maximizing query for token $i$; these $q_1,\ldots,q_n$ need not be the same.

Define

For each $i$, applying Jensen’s inequality (Lemma 9, with each $x_j$ set to $e^{\langle k_j,q_i\rangle/\sqrt{d}}$) gives

By the definition of $q_i$, $\sigma_i=e^{\langle k_i,q_i\rangle/\sqrt{d}}\big/\sum_{j=1}^{n}e^{\langle k_j,q_i\rangle/\sqrt{d}}$. Combining this with $\sigma_i/\lambda_i=\sum_{j=1}^{n}\sigma_j$, we have

Applying Equation (18) to the right-hand side,

We bound this expression in two ways. First, $\sum_i\lambda_i\langle k_i,q_i\rangle/\sqrt{d}\leq h$. Also, $\langle\bar{k},\bar{q}\rangle/\sqrt{d}=\sum_{i,j}\lambda_i\lambda_j\langle k_i,q_j\rangle/\sqrt{d}\geq\ell$, since every $q_j$ belongs to $\mathcal{Q}$. Hence

Second, centering the vectors (Lemma 12) gives $\sum_i\lambda_i\lVert k_i-\bar{k}\rVert^2=\sum_i\lambda_i\lVert k_i\rVert^2-\lVert\bar{k}\rVert^2\leq R^2\sqrt{d}-\lVert\bar{k}\rVert^2$, and the same bound holds for the queries. Applying Cauchy–Schwarz (Lemma 10) to (19), we obtain

The last two inequalities follow from the arithmetic-geometric mean inequality (Lemma 11) and Young’s inequality for inner products, $2\langle\bar{k},\bar{q}\rangle\geq-\lVert\bar{k}\rVert^2-\lVert\bar{q}\rVert^2$, respectively. Averaging this bound with $\log(\sum_{i=1}^{n}\sigma_i)\leq h-\langle\bar{k},\bar{q}\rangle/\sqrt{d}$ from (20) gives

Together with $\log(\sum_{i=1}^{n}\sigma_i)\leq h-\ell$, this proves the lemma.

For sensitivities $\sigma_i$ defined in Equation (12) with the supremum taken over $q \in\mathcal{Q}$, we can compute sampling scores $s_i$ satisfying $s_i \ge\sigma_i$ for every $i \in[n]$, and

Moreover, such sampling scores can be obtained by solving $n \cdot\max\{1,\lceil R^2\rceil\}$ convex optimization problems, using only the keys and the parameters $r_k$, $r_q$, $a$, $b$.

Proof. We group keys and queries by norm, compute sampling scores within each pair of groups by convex optimization, and then bound the sum of these scores using Lemma 4.

Partitioning by norm. Set $M = \max\{1,\lceil R^2\rceil\}$. Partition the keys into norm bins

Similarly, partition the possible query norms into $M$ intervals of width $r_q/M$, assigning each endpoint to the interval on its left and including zero in the first interval.

For $u,v \in[M]$, denote

If $i \in I_u$ and an admissible query $q$ has norm in the $v$-th interval, then

Thus, for each nonempty $I_u$, the compact convex domain

contains every admissible query whose norm lies in the $v$-th interval.

Bounding sensitivities locally. For $i \in I_u$, define the local sensitivity

with $\sigma_i^{uv}=0$ when $\widetilde{\mathcal{Q}}_{uv}$ is empty. For an admissible query $q$ whose norm lies in the $v$-th interval,

Consequently, $\sigma_i \le\sum_{v=1}^{M}\sigma_i^{uv}$.

Computing the sampling scores. For $i \in I_u$, define

By Lemma 13, this objective is convex and continuously differentiable, with an explicit gradient. Whenever $\widetilde{\mathcal{Q}}_{uv}$ is nonempty, it is the logarithm of the reciprocal of token $i$’s attention weight within $I_u$, so

Moreover, for every $z \in\widetilde{\mathcal{Q}}_{uv}$, $0 \le f_i^u(z) \le\log|I_u| + h_{uv} - \ell_{uv}$. The lower bound follows from the term $j=i$, and the upper bound follows from the defining inner-product bounds of $\widetilde{\mathcal{Q}}_{uv}$.

By Lemma 14 and the definition of $\widetilde{Q}_{uv}$, each nonempty domain is compact and convex. Applying Lemma 15 with $\eta= \log2$, we obtain $z_i^{uv} \in\widetilde{Q}_{uv}$ satisfying

Set $s_i^{uv} := 2e^{-f_i^{u}(z_i^{uv})}$. Exponentiating Equation (21) gives $\sigma_i^{uv} \le s_i^{uv} \le2\sigma_i^{uv}$. For empty domains, set $s_i^{uv} := 0$, so the same inequalities hold under the convention $\sigma_i^{uv}=0$. We are now ready to describe the sampling scores $s_i$. For $i \in I_u$, define

Together with the local sensitivity bound, this gives $\sigma_i \le\sum_{v=1}^{M}\sigma_i^{uv} \le s_i \le2\sum_{v=1}^{M}\sigma_i^{uv}$. Since $0 \in Q$, we have $\sigma_i \ge1/n$, so the scores are positive.

Bounding the sum within each pair of bins. Write $D_{uv}=B_{uv}-A_{uv}=\frac{R^2(u+v-1)}{M^2}\le\frac{2R^2}{M}\le2$. For $x_+=\max\{x,0\}$, the endpoint definitions give

Since $a_+ + (-b)_+ \le1$,

The keys in $I_u$ have norm at most $ur_k/M$, and queries in $\widetilde{Q}_{uv}$ have norm at most $vr_q/M$. The product of these bounds, divided by $\sqrt{d}$, is $B_{uv}$. For each nonempty domain, Lemma 4 therefore gives

The same bound holds trivially for empty domains. Hence

Summing over the bins. If $R^2<1$, then $M=1$, and the desired bound follows immediately. Otherwise, $M=\lceil R^2\rceil\le2R^2$. Since $u,v\le M$,

Using $R^2/(2M)\ge1/4$, we obtain

Finally, $0<\kappa\le1$ and $1-e^{-\kappa/4}\ge\kappa/8$, so

Each key participates in at most $M$ optimization problems, giving at most $nM$ such problems in total. ∎

Proofs from Section 3.2

In this section, we give the proof of Theorem 4, the lower bound in the importance sampling regime.

The construction and analysis are essentially the same as that in the uniform sampling regime § 2.2 with slightly more complicated parameter settings in order to capture all possible values of $a,b$.

Goal. We construct keys and queries such that each matching pair $(k_i,q_i)$ has scaled inner product $bR^2$, exceeding that of every nonmatching pair $(k_j,q_i)$, $j \ne i$, by at least $(1-o(1))\min\{b-a,(1+b)/2\}R^2$, where $R^2 = r_kr_q/\sqrt{d}$. We will then use this gap to make each query concentrate its attention on copies of its matching key, forcing an accurate summary to retain many key groups.

Construction. Let $e_0,e_K,e_Q$ be the first three standard basis vectors. Choose $m$ unit vectors satisfying Lemma 2 in dimension $D=d-3$, with pairwise inner products bounded in magnitude by $\xi=\sqrt{\frac{6\log(2n)}{D}}$. Embed them into the last $d-3$ coordinates of $\mathbb{R}^d$, and denote the resulting vectors by $u_1,\ldots,u_m$. Define

where we set

Note that $\alpha\ge0$ as $b \ge-1$ and $b \ge a$. Note that $\alpha+|\beta| \le1$ as $\alpha+\beta=2\alpha-b \le2\left(\frac{1+b}{2}\right)-b=1$ and $\alpha-\beta=b \le1$. Therefore, $1-\alpha-|\beta| \ge0$ and the keys and queries are well defined with norms exactly equal to $\lVert k_i\rVert=r_k$ and $\lVert q_i\rVert=r_q$. The components along $u_i$ and $e_0$ serve the purpose of achieving the desired inner products while the padding along $e_K$ and $e_Q$ increases the norms of the keys/queries to fix the denominator of cosine similarity.

Verifying the inner products. It remains to bound the scaled inner products between queries and keys. For any $i \in[m]$, the inner product between the matching pair is:

For any non-matching pair $i\ne j$,

The second choice in the minimum defining $\alpha$ ensures $b-\alpha(1+\xi)\ge a$. Thus, every key-query cosine similarity belongs to $[a,b]$, since $R^2\sqrt{d}=r_kr_q=\lVert k_i\rVert\lVert q_i\rVert$.

For any $i \ne j$, the scaled inner product between key-query pairs at the same index exceeds that of any other key-query pair by

By the definition of $\alpha$,

The last inequality uses $(1-x)/(1+x)\ge1-2x$ for $x\ge0$.

Constructing the value vectors. Apply Lemma 2 a second time to obtain $m$ unit value vectors $v_1,\ldots,v_m\in\mathbb{R}^d$ satisfying

Since $m\le n$, the assumption $d=\omega(\log n)$ ensures that this is possible.

Full construction

Proof of Theorem 4. Choose $t=\lfloor n/m\rfloor\ge1$ and $0\le s<m$ such that $n=tm+s$. For $i\in[m]$, let $N_i=t+1$ if $i\le s$ and $N_i=t$ if $i>s$. The attention instance $K,V$ will have $N_i$ copies of the key-value pair $(k_i,v_i)$. By the setting of $t,s,N_i$, there will be exactly $n$ key-value pairs. As $t\ge1$, $N_j/N_i\le2$ for all $i,j\in[m]$.

At distinct query $q_i$, let $\alpha_i$ be the attention mass on the copies of $(k_i,v_i)$. By Equation (25),

Recall the condition $m\le e^{(1-2\xi)\min\{b-a,(1+b)/2\}R^2/32}$. Then,

So, attention focuses significant mass on $v_i$:

The first inequality follows because all value vectors have norm one, so their distance from $v_i$ is at most two.

Let $H_i$ be the event that the coreset retains at least one copy of $(k_i,v_i)$. On the complementary event, the approximation at $q_i$ is a convex combination of the value vectors $v_j$ with $j\ne i$. This follows directly from the nonnegative weights in the definition of query-oblivious weighted token-eviction. Therefore

as all distinct value vectors are nearly orthogonal.

On the other hand, Equation (27) gives $\langle v_i,\operatorname{Attn}(K,V,q_i)\rangle\ge1-1/8$. Projecting the error onto the unit vector $v_i$ yields, whenever $H_i$ does not hold,

Hence success at the fixed query $q_i$ implies $H_i$.

The per-query guarantee of a query-oblivious weighted token-eviction policy now gives $\Pr[H_i]\ge2/3$ for each $i$ separately. Since each retained pair represents at most one group,

By definition, $\mathrm{OPT}_{\mathrm{Attn}}$ is the optimal worst-case number of tokens retained by an algorithm, so, from above, $\mathrm{OPT}_{\mathrm{Attn}}\ge2m/3$.

Experiments

Q/K Cosine similarity

Vector (PDF)

Vector (PDF)

Key/Query norms