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

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

cs.LGOct 26, 2026

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

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

Organizations: Massachusetts Institute of Technology (MIT). · École Polytechnique Fédérale de Lausanne (EPFL).

Abstract

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.

Figures & tables

Figure 1: Cosine similarities between the final 256 queries and preceding context keys in layer 8 of Qwen3-8B, for an 8K-token LongBench v2 input, excluding the first key (attention sink).

Model	Layer 4	Layer 8	Layer 16	Layer 24
Llama 3.1 8B	$-0.584\pm 0.148$	$-0.725\pm 0.098$	$-0.681\pm 0.076$	$-0.688\pm 0.105$
Qwen3 8B	$0.068\pm 0.037$	$-0.146\pm 0.034$	$0.024\pm 0.044$	$-0.143\pm 0.080$
Ministral 3 8B	$-0.724\pm 0.047$	$-0.628\pm 0.086$	$-0.613\pm 0.078$	$-0.658\pm 0.045$

Table 1: Cosine similarity between mean query and key vectors. Values are means

\pm

standard deviations across example/context-length configurations common to all models and layers.

Reference	Model / regime	Upper bound	Lower bound
Kochetkova et al. (2025)	Coreset size (tokens)	$\widetilde{O}\!\left(\sqrt{d}e^{2R^{2}}\right)$	$\widetilde{\Omega}\!\left(1\right)$
Haris and Onak (2025) a	Space complexity (bits)	$\widetilde{O}(nd)$	$\Omega(nd)$
Liberty et al. (2026) b	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. (2026) c	Space complexity (bits)	$\widetilde{O}\!\left(e^{(1+o(1))R^{2}}\right)$	$\widetilde{\Omega}\!\left(e^{(1-o(1))R^{2}}\right)$
This work d , $b-a\leq\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 work d , $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)$

Table 2: Worst-case space complexity of attention.

\widetilde{O},\widetilde{\Omega}

suppress

\operatorname{poly}(\varepsilon^{-1}\log n)

factors.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Query–key cosine similarities across models and attention layers. Rows show Llama 3.1 8B, Qwen3 8B, Ministral 3 8B; columns show layers 4, 8, 16, 24. All panels use the same prepared LongBench v2 example. Each panel uses an input of 8,192 tokens and compares the final 256 queries from query head 0 against 7,935 earlier-context keys from the corresponding KV head 0. The initial key is excluded. Cosines are computed after model normalization and RoPE.

Figure 7: Query–key cosine similarities across inputs. The plot uses Qwen3 8B at layer 8 on (a) LongBench v2, (b) InfiniteBench, (c) SCBench. Each panel uses an input of 8,192 tokens and compares the final 256 queries from query head 0 against 7,935 earlier-context keys from the corresponding KV head 0. The initial key (attention sink) is excluded. Cosines are computed after model normalization and RoPE.

Figure 8: Query–key cosine similarities with initial token (attention sink) Qwen3 8B at layer 8 on LongBench v2 example input of 8,192 tokens and compares the final 256 queries from query head 0 against a sample of 16 earlier-context keys from the corresponding KV head 0 and the initial key (attention sink). Cosines are computed after model normalization and RoPE.

Figure 9: Key and query magnitudes across models. Layer 8 on the LongBench v2 input (8,192 tokens). Columns show Llama 3.1 8B, Qwen3 8B, Ministral 3 8B. The top row uses 7,935 keys per model: all eligible earlier-context keys. Blue baseline marks indicate the initial key (attention sink) norm in each model. The bottom row uses the final 256 queries.

Attention Approximation

See all 19 papers

Towards Tight Bounds for Streaming Attention
Jun 5, 2026Justin Y. Chen, Ying Feng, Piotr Indyk +3Efficient Attention KV-Cache Compression
MIT · EPFL
Self-Pruned Key-Value Attention: Learning When to Write by Predicting Future Utility
May 13, 2026Gergely Szilvasy, Manuel Faysse, Maria Lomeli +5Self-Attention KV Caching
Meta FAIR · MICS, CentraleSupélec
Thin Keys, Full Values: Reducing KV Cache via Low-Dimensional Attention Selection
Feb 16, 2026Hengshuai Yao, Xing Chen, Ahmed Murtadha +1Memory-Efficient Attention KV-Cache Compression
Sapient Intelligence · Department of Computing Science, University of Alberta

Jun 5, 2026J/K move · Enter open · S save

Towards Tight Bounds for Streaming Attention

Justin Y. Chen, Ying Feng, Piotr Indyk, Michael Kapralov+2

MIT · EPFL

Efficient Attention KV-Cache Compression Information-Theoretic Lower Bounds Streaming Algorithms

Open paper AskPDF arXiv

The attention mechanism is a cornerstone of modern transformer architectures. However, its expressive power comes at the cost of quadratic runtime and linear space usage. In particular, the classical transformer architecture explicitly stores all previously seen input elements (tokens) in order to generate the next one. The problem of implementing a transformer in limited space, known as KV cache compression, has received much interest over the past few years, spurring the development of powerful heuristics. Recent works of Haris et al, COLT'25 and Kochetkova et al, NeurIPS'25, formalized KV cache compression as the streaming attention approximation problem, providing both upper bounds (based on discrepancy theory) and information theoretic lower bounds. However, those papers left open a significant gap between the upper and lower bounds. For example, the space usage of their algorithms increases with the precision parameter, but the lower bound does not get stronger. In this work, we revisit the streaming attention approximation problem and provide nearly tight bounds on its space complexity. On the algorithmic side, we achieve the result through a surprisingly tight interplay between three distinct methods for kernel density estimation: discrepancy-based coreset constructions (e.g., Charikar-Kapralov-Waingarten'24), the polynomial method (e.g., Greengard-Rokhlin'87, Alman-Song'23), and space partitioning (e.g., Andoni-Laarhoven-Razenshteyn-Waingarten'17, Charikar-Kapralov-Nouri-Siminelakis'20). On the lower bound side, our main technical contribution is a new technique for using the INDEX problem with a large amount of side information that we hope will prove useful in other high dimensional geometric estimation problems.

Abstract

Attention Approximation

Explore similar work

Towards Tight Bounds for Streaming Attention

Self-Pruned Key-Value Attention: Learning When to Write by Predicting Future Utility

Thin Keys, Full Values: Reducing KV Cache via Low-Dimensional Attention Selection