PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Justin Y. Chen, Ying Feng, Piotr Indyk, Michael Kapralov, Ekaterina Kochetkova, Boris Prokhorov

arXiv:2610.35020v1cs.LGcs.DS

TL;DR

KV-cache compression seeks to reduce the memory needed to store context while approximating attention, but nearly tight worst-case bounds permit key-query alignment unlike the small similarities observed in practice. This paper studies attention approximation under bounded key-query similarity and gives a nearly space-optimal algorithm and matching lower bound. Its lower-bound distribution exhibits anisotropy, and uniform sampling is optimal in a parameter regime.

  • Problem

    KV-cache compression addresses the memory cost of storing all prior tokens during decoding, while existing nearly tight worst-case attention bounds allow key-query alignment unlike the small similarities observed in practice.

    Source: §1 · p. 1, §1.1 · p. 2

  • Method

    The paper studies query-oblivious weighted token eviction for attention approximation under bounded key-query similarities and establishes upper and lower space bounds.

    Source: §1 · p. 1, §1.1 · p. 3

  • Results

    Uniform sampling achieves asymptotically tight space bounds in the uniform regime, and importance sampling achieves asymptotically tight bounds in the importance regime.

    Source: §1.1 · p. 3

  • Takeaways & Limitations

    Uniform sampling is optimal for a wide parameter range, and the lower-bound construction exhibits the anisotropy observed in key-query geometry.

    Source: §1.1 · pp. 2–3

    Real KV caches may have additional structure that makes them easier to compress than the paper's hard instances.

    Source: §4 · p. 9

Abstract

from arXiv · show

A growing body of work on KV-cache compression aims to reduce the memory requirements of autoregressive generation. These methods compress the stored context while approximately preserving attention computation, enabling generation with reduced memory consumption. Almost tight theoretical results on the worst case space complexity of attention have recently been obtained, but they are achieved by instances where queries can be perfectly aligned with keys. This directly contradicts the observed anisotropy of the key/query geometry in real transformers. In particular, it has been shown empirically that key/query cosine similarities are small even for most correlated pairs, i.e., near-perfect alignment of queries with keys does not happen in practice. In this paper, we study the space complexity of attention approximation under the assumption that key/query cosine similarities belong to a small interval around zero, i.e., a query cannot align too much with any key. We give a nearly space optimal algorithm and a corresponding space lower bound. Surprisingly, the distribution of key/query cosine similarities in our lower bound exhibits the anisotropy property empirically established and studied in the literature. Additionally, we show that the simple uniform sampling algorithm is nearly optimal for a typical range of parameters, supporting its surprising effectiveness on real datasets.

1 INTRODUCTION

KV-cache compression addresses the memory burden of storing all tokens during transformer decoding by approximating attention with a smaller representation. This section frames query-oblivious weighted token eviction and contrasts nearly tight worst-case bounds with practical key-query geometry.

1 INTRODUCTION

  • Decoding requires storing all previously seen tokens, making high LLM memory consumption a deployment challenge that KV-cache compression aims to address.Compression methods approximate attention using small space, either by quantizing vectors or evicting tokens.Source: §1 · p. 1
  • Query-oblivious token eviction selects a random subset of token indices and positive weights before seeing the specific query, evicting the remaining tokens.The resulting weighted summary is used to approximate attention for a given query.Source: §1 · p. 1
  • Success is required for each fixed query independently of the summary randomness, not simultaneously for all queries; key, query, and value norms are bounded.The stated assumptions bound key norms by r_k, query norms by r_q, and value norms by 1.Source: §1 · p. 2
  • Prior work established nearly tight worst-case space bounds, but those results allow queries to align with keys, unlike practical instances highlighted here.The cited theoretical dependence is sufficient and necessary in the worst case, while the paper notes a divergence from practical instances.Source: §1 · p. 2

1. The space lower bound of erkrq/

The paper studies attention approximation when key-query cosine similarities are restricted to a small interval around zero, contrasting with worst-case instances that permit near-perfect alignment. It gives nearly tight sampling bounds and connects the uniform-sampling regime to empirical similarity ranges and anisotropic key-query geometry.

1. The space lower bound of erkrq/

  • The analysis assumes every key-query cosine similarity lies in a small interval around zero, so queries cannot align significantly with keys.This differs from prior worst-case instances where queries can align perfectly with a key.Source: §1.1 · p. 2, §1 · p. 2

1.1 OUR RESULTS

  • Uniform sampling has asymptotically tight space bounds in its regime, while importance sampling has asymptotically tight bounds in the other regime.The paper places widely used models in or near the uniform-sampling regime.Source: Figure 2 · p. 3, §1.1 · p. 3
  • Uniform sampling is optimal when a < 0 < b and b − a ≤ 1 + b/2; importance sampling applies for larger cosine similarities.The empirical (a, b) values for widely used models lie in or near the uniform-sampling regime.Source: §1.1 · pp. 2–3
  • The lower-bound construction is anisotropic: keys and queries are shifted in opposite directions along a shared direction.This resembles anisotropy observed empirically in practice.Source: §1.1 · p. 3

1.2 RELATED WORK

  • The retained-token lower bounds do not apply to arbitrary quantized or low-rank KV representations, which fall outside the weighted-subset model.The related work also distinguishes cache-selection methods based on historical attention or prompt observations from this explicit similarity-restricted analysis.Source: §1.2 · p. 4

2 UNIFORM SAMPLING REGIME

In the small-cosine regime, uniform sampling approximates attention, and a matching lower bound shows that many tokens can still be necessary. The lower-bound construction also exhibits geometric features observed in real query-key data.

2.1 UPPER BOUND

  • Uniform sampling gives an additive-ε attention approximation with probability at least 2/3 for each fixed admissible query in the stated small-cosine regime.The regime has −1 ≤ a < 0 < b ≤ 1 and b − a ≤ 1 + b.Source: §2.1 · pp. 4–5

2.2 LOWER BOUND

  • For this dataset, which has m groups of identical key-value pairs each paired with a different query, satisfying Equation (3) at each of these queries requires OPTAttn ≥2m/3.Source: §2.2 · p. 5

2.2.1 LOWER BOUND CONSTRUCTION

  • Keys and queries are shifted oppositely along e0, while each pair shares a nearly orthogonal direction, making its matching key dominate attention.Nearly orthogonal value vectors ensure that other groups’ values cannot well-approximate the attention output.Source: §2.2 · p. 6

2.2 LOWER BOUND

  • The hard instance shares observed query-key geometry: separation along a common direction and otherwise similar distributions at different scales.The empirical visualization is for Llama 3.1 8B on a LongBench v2 example.Source: §2.2 · p. 6, Figure 5 · p. 6

3 IMPORTANCE SAMPLING REGIME

The upper bound uses sensitivity-based importance sampling to construct a weighted attention summary, with uniform sampling as a special effective regime. The lower-bound analysis reuses the uniform-regime construction with a changed condition on m.

3.1 UPPER BOUND

  • Importance sampling assigns higher probabilities to tokens with greater possible attention contribution, using sensitivity upper bounds as sampling scores.The method bounds total sensitivity by partitioning key and query-norm ranges and using smaller convex relaxations to compute scores.Source: §3.1 · pp. 7–9
  • With fixed a and b, the importance-sampling scores yield an additive-ε attention approximation with probability at least 2/3 for each fixed admissible query.The score sum is bounded by S ≤ 128e²κ⁻²e^κR², where κ = min{b − a, (1 + b)/2}.Source: §3.1 · p. 9

3.2 LOWER BOUND

  • The importance-regime lower bound reuses the uniform-regime construction and analysis, changing the theorem’s condition on m.The full theorem statement and proof are deferred to Section E.Source: §3.2 · p. 9

4 OPEN PROBLEMS

The open problems ask how transformer training produces the observed key-query geometry and whether real KV caches contain additional structure that makes them easier to compress.

4 OPEN PROBLEMS

  • The paper asks what assumptions on data and training dynamics would explain small key-query cosine similarities and negative alignment between mean keys and queries.It identifies a theoretical explanation for how transformers learn this geometry as an open question.Source: §4 · p. 9
  • The paper asks whether real KV caches have additional structure beyond the hard instances that could make them easier to compress.The hard instances share some geometric features with trained models, but the passage leaves further real-cache structure unresolved.Source: §4 · p. 9

B PROOFS FROM SECTION 2.1

The proof establishes the uniform-sampling approximation guarantee by controlling the attention denominator and the error numerator separately.

B PROOFS FROM SECTION 2.1

  • Uniform sampling gives an additive-ε attention approximation with probability at least 2/3 for each fixed admissible query.The guarantee applies to the stated uniform sampling procedure.Source: B PROOFS FROM SECTION 2.1 · p. 14
  • The proof uses Chebyshev’s inequality to control the denominator and Markov’s inequality to bound the error numerator.Together, the bounds imply the stated approximation guarantee.Source: B PROOFS FROM SECTION 2.1 · p. 15

C PROOFS FROM SECTION 2.2

The proof establishes that high-dimensional space contains large sets of unit vectors with pairwise inner products bounded near zero. It uses random sign vectors and a probability bound to show such a set exists.

C PROOFS FROM SECTION 2.2

  • For M ≥ 2 and 0 < ξ < 1, dimension D ≥ 6 log(2M)/ξ^2 suffices for M unit vectors whose pairwise inner products have magnitude at most ξ.The proof samples coordinates independently from {−D^−1/2, D^−1/2}; Hoeffding’s inequality and a union bound make the probability of any violating pair less than one.Source: C PROOFS FROM SECTION 2.2 · p. 15

C.1 PROOF OF THEOREM 3

The lower-bound proof groups repeated key-value pairs and constructs queries that concentrate attention on their matching groups. Accurate answers across the queries therefore require retaining many groups, yielding a worst-case token lower bound.

C.1 PROOF OF THEOREM 3

  • The construction divides n key-value pairs among m groups, giving each group either ⌊n/m⌋ or ⌊n/m⌋+1 identical copies.The group sizes differ by at most one, so the ratio between any two group sizes is at most 2.Source: §C.1 · p. 16
  • Without any copy of a query’s matching pair, the coreset’s attention output differs from the true output by at least 3/4, exceeding ε.The matching group carries attention mass at least 15/16, while the unit value vectors have pairwise inner products bounded in magnitude by 1/8.Source: §C.1 · p. 16
  • OPTAttn ≥ 2m/3 for the constructed instance, so any algorithm’s optimal worst-case retained-token count grows with the number of groups.Each accurately answered query requires a copy of its matching pair, and each retained pair represents at most one group.Source: §C.1 · p. 16

D PROOFS FROM SECTION 3.1

The proof analyzes Algorithm 1’s weighted attention summary using sensitivity upper bounds and a sampling-based estimator. It also describes how to compute usable sampling scores through norm partitioning and convex optimization.

D PROOFS FROM SECTION 3.1

  • Under ∥v_i∥ ≤ 1 and scores s_i ≥ σ_i, Algorithm 1’s estimator has an approximation guarantee and the algorithm retains at most min{n,m} tokens.The probability in the estimator analysis is over the algorithm’s sampling.Source: D PROOFS FROM SECTION 3.1 · pp. 16–17
  • The sampled numerator and denominator are unbiased, and concentration bounds on them yield the estimator analysis.The proof applies Chebyshev’s inequality to the denominator and Markov’s inequality to the centered numerator, then combines the bounds with a union bound.Source: D PROOFS FROM SECTION 3.1 · p. 17
  • Sampling scores that upper-bound every sensitivity can be computed using n · max{1, ⌈R^2⌉} convex optimization problems.The computation uses the keys and parameters r_k, r_q, a, and b; the proof groups keys and queries by norm to bound the scores.Source: D PROOFS FROM SECTION 3.1 · p. 19
  • The score computation partitions key norms and query norms into bins, then bounds local sensitivities within each pair of bins.For each key bin, the proof uses compact convex query domains and optimizes a convex objective related to the reciprocal of a token’s attention weight.Source: D PROOFS FROM SECTION 3.1 · p. 19

E PROOFS FROM SECTION 3.2

The importance-sampling lower-bound proof constructs a fixed dataset and queries for which accurate attention summaries must retain many tokens. Its construction keeps key-query cosine similarities within the prescribed interval while making matching pairs more influential.

E PROOFS FROM SECTION 3.2

  • For the constructed dataset, satisfying the attention guarantee at each of m queries requires OPTAttn ≥ 2m/3.The construction uses n key-value pairs with unit-norm values and keys of norm r_k.Source: E PROOFS FROM SECTION 3.2 · p. 21
  • Nearly orthogonal vectors embedded in dimension d−3 provide the construction’s group directions, with dimension d = ω(log n) ensuring their availability.The construction adds norm padding along fixed directions so keys and queries have norms r_k and r_q.Source: E PROOFS FROM SECTION 3.2 · pp. 21–22
  • The construction places every key-query cosine similarity in [a,b] while giving each matching pair a larger scaled inner product than nonmatching pairs.The matching-pair gap is at least (1−o(1)) min{b−a,(1+b)/2}R^2; the proof uses that gap to concentrate attention on matching key groups.Source: E PROOFS FROM SECTION 3.2 · p. 21

E.1 FULL CONSTRUCTION

The construction spreads n key-value pairs across m groups and gives each group a query whose attention concentrates on its value. A successful approximation at each query must retain that group, yielding a lower bound of 2m/3 retained tokens.

E.1 FULL CONSTRUCTION

  • The construction creates exactly n key-value pairs by assigning each of m groups either ⌊n/m⌋ or ⌊n/m⌋ + 1 copies.The group sizes differ by at most one, so N_j/N_i ≤ 2 for all groups.Source: §E.1 · p. 22
  • At query q_i, attention concentrates significant mass on v_i, while the distinct value vectors are nearly orthogonal.The construction bounds the competing key-query similarities so the query’s own group receives substantial attention.Source: §E.1 · p. 22
  • The per-query guarantee gives Pr[H_i] ≥ 2/3, and the group-retention argument yields OPTAttn ≥ 2m/3.Each retained pair represents at most one group, and success at query q_i implies retaining a copy of its group’s pair.Source: §E.1 · p. 22

F EXPERIMENTS

The experiments describe cosine-similarity comparisons across models, layers, and datasets, plus key and query magnitudes. The supplied passages specify these figure setups but report no observed outcomes.

F.1 Q/K COSINE SIMILARITY

  • Figure 6 compares query–key cosine similarities across three models and four layers using the same 8,192-token LongBench v2 example.Each panel compares the final 256 queries with 7,935 earlier-context keys, excluding the initial key.Source: Figure 6 · p. 23
  • Figure 7 compares cosine similarities across LongBench v2, InfiniteBench, and SCBench for Qwen3 8B at layer 8.The setup excludes the initial attention-sink key; Figure 8 instead includes it alongside a sample of 16 earlier-context keys.Source: Figure 7 · p. 23, Figure 8 · p. 24

F.2 KEY/QUERY NORMS

  • Figure 9 presents key magnitudes in its top row and query magnitudes in its bottom row for three models at layer 8.The input is an 8,192-token LongBench v2 example, and blue baseline marks indicate each model’s initial-key norm.Source: Figure 9 · p. 24

Related papers

Loading 2610.35020v1…