🤖 AI

Nearly Optimal Attention Coresets

This paper establishes the existence of nearly optimal-sized attention coresets for unit-norm keys and values, providing an improved upper bound of $O(\sqrt{d} e^{\rho+o(\rho)}/\varepsilon)$ and a matching lower bound of $\Omega(\sqrt{d} e^{\rho}/\varepsilon)$ that outperforms previous results.

Original authors: Edo Liberty, Alexandr Andoni, Eldar Kleiner

Published 2026-05-08

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Edo Liberty, Alexandr Andoni, Eldar Kleiner

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are running a massive library (a modern AI model) where every book has a "Key" (a summary of what it's about) and a "Value" (the actual content). When a reader asks a question (a "Query"), the librarian uses a special mechanism called Attention to scan all the books, figure out which ones are most relevant, and summarize their content into a single answer.

As the library grows to hold millions of books, the librarian's desk becomes cluttered. Keeping track of every single book's Key and Value takes up too much space and slows everything down. The goal of this paper is to figure out: How many books can we throw away from the desk while still giving the reader the exact same answer?

Here is the breakdown of the paper's findings using simple analogies:

1. The Problem: The "Cherry-Picker"

The authors explain that you can't just throw away random books. If a reader asks a very specific, intense question (like "Find me the one book that mentions a specific rare word"), the librarian needs to be able to isolate that single book perfectly. If you remove too many books, you might lose that specific one, and the answer will be wrong.

In technical terms, if a reader's question is allowed to be infinitely "loud" or intense, you can't compress the library at all. You'd have to keep every single book.

The Solution: The paper says, "Let's agree that readers won't shout too loudly." If we limit how intense the questions can be (a "bounded norm"), we can safely throw away most of the books and keep only a tiny, carefully selected group that represents the whole library.

2. The Magic Trick: The "Balancing Act"

The core of the paper is a mathematical method to pick which books to keep. The authors use a technique called Coreset Selection.

Imagine you have a giant pile of weights (the books) on a scale. You want to remove half the weights but keep the scale perfectly balanced so it doesn't tip.

The Old Way: Previous methods tried to balance the scale by looking at the weights one by one, which was slow and left a lot of extra "noise" (error).
The New Way: The authors use a clever mathematical trick (based on a theorem called Banaszczyk's vector balancing). They imagine the weights as arrows pointing in different directions. They assign a "plus" or "minus" sign to every book.
- If the signs are chosen perfectly, the "plus" books and "minus" books cancel each other out almost completely.
- The books with the "plus" signs become your new, tiny library.
- Because the "minus" books cancelled out the noise, the "plus" books still represent the whole group perfectly.

3. The Result: "Nearly Optimal" Size

The paper proves two main things:

The Good News (Upper Bound): They found a way to shrink the library down to a size of roughly $\sqrt{d} \times e^{\rho}$ (where $d$ is how complex the books are, and $\rho$ is how loud the questions can be). This is the smallest size they could mathematically prove is possible using their method. It's much smaller than what anyone had found before.
The Bad News (Lower Bound): They also proved that you cannot go much smaller than this. If you try to shrink the library further, there will inevitably be some questions where the answer becomes wrong.

Think of it like packing a suitcase. The authors found a way to fold your clothes so tightly that the suitcase is almost as small as physically possible. They also proved that you can't fold them any tighter without crushing the clothes.

4. Why This Matters

In the world of AI, "Keys" and "Values" are the memory of the model. As AI models try to remember longer and longer conversations (context), this memory becomes huge and expensive.

This paper provides a theoretical guarantee that we can compress this memory significantly without losing accuracy, as long as the questions aren't too extreme. It tells engineers: "You don't need to keep 100% of the data. You can keep a tiny fraction, and mathematically, the AI will still work just as well."

Summary in a Sentence

The authors discovered a mathematical "folding technique" that allows AI models to shrink their memory down to the smallest possible size without losing accuracy, proving that this new size is almost the absolute limit of what is physically possible.

Based on the provided paper, here is a detailed technical summary of "Nearly Optimal Attention Coresets."

Problem Definition

The paper addresses the challenge of compressing the Key-Value (KV) cache in transformer-based generative models. As context windows expand, the size of the KV-cache becomes a primary bottleneck for deploying large language models. The authors formalize the problem of Attention Coreset Selection: given a set of $n$ key-value pairs $(K, V)$ in $\mathbb{R}^d$ (where keys and values are unit-norm), the goal is to find a small subset $(K', V')$ of size $S$ such that the attention output is approximated within an error $\epsilon$ for all query vectors $q$ with a bounded norm $\|q\| \le \rho$ .

The attention mechanism is defined as:
$\text{Attn}(q, K, V) = \frac{A(q, K, V)}{B(q, K)} = \frac{\sum_i e^{q^T k_i} v_i}{\sum_i e^{q^T k_i}}$

The authors note that without bounding the query norm $\rho$ , lossy compression is impossible. A query with a sufficiently large norm can "cherry-pick" a single key-value pair, forcing the sketch to store the entire dataset to maintain accuracy. Therefore, the problem is constrained to queries where $\|q\| \le \rho$ .

Methodology

The core technical approach relies on discrepancy theory, specifically leveraging Banaszczyk's vector balancing theorem. The methodology involves three main components:

Halving via Discrepancy: The authors employ a "halving trick" where they assign signs $\sigma_i \in \{-1, 1\}$ to the terms in the attention sum. By selecting signs to minimize the discrepancy (the error term $E(q) = \sum \sigma_i e^{q^T k_i} v_i$ ), they can construct a subset $(K', V')$ containing the terms where $\sigma_i = 1$ . This subset is at most half the size of the original set, and the error is bounded by the discrepancy of the signed sum.
Tensor Norm Balancing: To bound the error for the vector-valued numerator $A(q, K, V)$ , the authors expand the exponential term $e^{q^T k_i}$ using a Taylor series. This transforms the problem into balancing a sum of tensor products ( $k_i^{\otimes m} \otimes v_i$ ). They define a specific scaling factor $z(m)$ and construct a high-dimensional vector $x_i$ containing these scaled tensor components.
Convex Body Construction: They apply Banaszczyk's theorem to show that there exist signs $\sigma_i$ such that the sum $\sum \sigma_i x_i$ lies within a specific convex body $K$ . This body is defined by constraints on the dual norms of the tensor sums. By carefully choosing the geometry of this convex body, they ensure the error remains small across all tensor orders $m$ .
Simultaneous Balancing: To handle the full attention mechanism (both the numerator $A$ and the denominator $B$ ), as well as the centering of keys, the authors prove a simultaneous balancing lemma. This allows them to satisfy the discrepancy bounds for multiple objectives (numerator, denominator, and key centering) simultaneously using a single set of signs, with only a constant factor increase in the error bound.
Iterative Compression: The halving process is applied iteratively ( $\ell$ times) until the accumulated error is within the desired tolerance $\epsilon$ . This reduces the dataset size from $n$ to $O(\sqrt{d} e^\zeta / \epsilon)$ .

Key Contributions and Results

1. Upper Bound (Theorem 1 & 4):
The paper establishes the existence of an attention coreset of size:
$S = O\left( \frac{\sqrt{d}}{\epsilon} e^{\zeta} \right)$
where $\zeta = \rho + \frac{1}{2}\log(\rho) + o(\log \log \rho)$ .
This result improves upon the previous best-known upper bound (from streaming attention approximation literature, e.g., [30]) by a factor of roughly $\log(n) \cdot e^{\rho - o(\rho)}$ . The authors achieve this by surgically modifying and simplifying existing discrepancy arguments to handle vector-valued functions more tightly.

2. Lower Bound (Theorem 2 & 7):
The authors prove a matching lower bound, demonstrating that any $\epsilon$ -coreset for unit-norm vectors and queries of norm at most $\rho$ must have a size of:
$S = \Omega\left( \frac{\sqrt{d}}{\epsilon} e^{\rho} \right)$
This lower bound is derived via a reduction from the one-way Indexing problem, utilizing spherical codes to construct hard instances where distinguishing between bits requires a large sketch size.

3. Near-Optimality:
Comparing the upper and lower bounds, the gap is a factor of $O(e^{\frac{1}{2}\log(\rho) + O(\log \log \rho)}) = O(\rho^{1/2 + o(1)})$ . The authors claim this establishes a "nearly optimal" coreset size, as the bounds match up to a sub-polynomial factor in $\rho$ .

Significance and Claims

The paper claims to resolve the question of the optimal size of coresets for attention mechanisms with error guarantees.

Theoretical Breakthrough: It provides the first tight existential bounds for attention coresets, moving beyond heuristic compression methods (like clustering or distillation) to provable guarantees based on discrepancy theory.
Practical Implication: The results suggest that for bounded-norm queries, the KV-cache can be compressed to a size that depends on the dimension $d$ , the error tolerance $\epsilon$ , and the query norm $\rho$ , but is independent of the total number of tokens $n$ (once $n$ is sufficiently large).
Limitations: The authors note that without the constraint on query norm $\rho$ , no compression is possible. Furthermore, they acknowledge that the remaining gap between the upper and lower bounds (roughly $\sqrt{\log n}$ in the worst case) likely requires fundamentally different approaches to close.

The paper does not propose new experimental implementations or specific architectural changes for transformers but focuses strictly on the theoretical limits of data compression for the attention mechanism.