🤖 machine learning

Towards Tight Bounds for Streaming Attention

This paper resolves the significant gap between existing upper and lower bounds for the streaming attention approximation problem by establishing nearly tight space complexity bounds through a novel combination of kernel density estimation techniques and a new lower bound method based on the INDEX problem with side information.

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

Published 2026-06-08

📖 6 min read🧠 Deep dive

CC BY 4.0

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

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 trying to build a super-smart robot that can read a book and then write a new chapter based on what it just read. To do this, the robot needs to remember every word it has read so far (the "context") and figure out which of those words are most important for the next sentence it wants to write.

In the world of AI, this process is called Attention. The problem is, as the book gets longer, the robot's memory gets clogged. It has to keep a giant list of every single word it's ever seen, which takes up a massive amount of space and slows everything down.

This paper is like a team of engineers who found a way to shrink that giant memory list down to a tiny, efficient size without losing the robot's ability to understand the story. They figured out the absolute best (or "tightest") way to do this, proving that you can't do much better than their method.

Here is how they did it, explained with some everyday analogies:

1. The Problem: The "Giant Library" vs. The "Pocket Note"

Think of the robot's memory as a library.

The Old Way: Every time the robot reads a new word, it puts a full, heavy encyclopedia on a shelf. If the book has 1,000 words, the robot needs 1,000 encyclopedias. This is slow and expensive.
The Goal: The robot wants to keep a "Pocket Note" instead. It wants to summarize the whole library into a few key sentences that still let it answer any question accurately.

Previous researchers tried to make these pocket notes, but they left a big gap between how small they could make the note and how small they actually made it. They didn't know the true limit.

2. The Solution: Three Tools for One Job

The authors of this paper realized that to shrink the memory perfectly, you need to use three different tools at the same time, depending on how "hot" or "cold" the data is (a concept they call "temperature").

Tool A: The "Moment" Sketch (The Snapshot)
Imagine you want to describe a crowd of people. Instead of listing every person, you take a photo that captures the average height, the average weight, and the general mood. This is a "sketch." It's great for describing the crowd when everyone is spread out and mixed up (the "high temperature" regime). The authors combined this with some advanced math (polynomials) to make the sketch incredibly efficient.
Tool B: The "Discrepancy" Filter (The Balanced Scale)
Sometimes, the crowd isn't mixed; maybe there's a group of tall people on the left and short people on the right. A simple photo doesn't work well here. Instead, you need a "filter" that balances the groups so you don't lose the difference. The authors used a mathematical trick called "discrepancy theory" to create a tiny group of people (a "coreset") that perfectly represents the balance of the whole crowd.
Tool C: The "Space Partition" Map (The Neighborhoods)
If the crowd is clustered into tight neighborhoods (like a "low temperature" regime where the robot is hyper-focused on just a few words), the authors realized you shouldn't treat the whole library as one big room. Instead, you should break the library into small rooms and summarize each room separately. They developed a way to find these clusters, move them to the center of the room (re-centering), and then shrink them down.

The Magic: The paper shows that by switching between these three tools depending on the situation, you can get a memory size that is almost as small as mathematically possible.

3. The "Tight" Result: No More Guessing

Before this paper, scientists were guessing how small the memory could get. They had a "best guess" for the smallest size (Upper Bound) and a "minimum possible" size (Lower Bound), but there was a huge gap between them.

The Analogy: Imagine you are trying to fit a suitcase into a car trunk. Previous researchers said, "It might fit if we squeeze it really hard," but they didn't know if the trunk was actually big enough.
This Paper: The authors measured the suitcase and the trunk with a laser ruler. They proved, "Yes, it fits, and here is the exact amount of space you need. You cannot fit it any smaller, and you don't need any more space than this."

They proved that for a wide range of scenarios, their method is nearly perfect. If you try to make the memory smaller than their method, the robot will start making mistakes. If you try to make it larger, you are just wasting space.

4. How They Proved It (The "Spy" Game)

To prove that you can't do better than their method, they used a clever trick involving a game of "20 Questions" (called the INDEX problem in math).

The Setup: Imagine a spy (Alice) has a secret code (a long string of 0s and 1s). She sends a tiny message to her partner (Bob). Bob needs to guess one specific bit of the code.
The Trick: The authors showed that if the robot's memory was smaller than their limit, the spy could use the robot's memory to send a message that was too small to solve the game. Since we know from math that the message must be a certain size to solve the game, the robot's memory must be at least that big.
The Innovation: They added a twist where the spy sends a little bit of "side information" (like a hint) to help Bob. This allowed them to prove the limit is even tighter than before, closing the gap that previous researchers couldn't fix.

Summary

In simple terms, this paper is a masterclass in compression.

The Problem: AI models are too hungry for memory.
The Fix: The authors built a new system that uses a mix of sketches, filters, and neighborhood maps to summarize data perfectly.
The Proof: They proved mathematically that this system is the best possible one. You can't shrink the memory any further without breaking the AI's brain.

They didn't just build a better tool; they drew the map showing exactly where the edge of the cliff is, so no one else has to waste time trying to walk off it.

Technical Summary: Towards Tight Bounds for Streaming Attention

Problem Statement

The paper addresses the Streaming Attention Approximation Problem, a fundamental challenge in deploying transformer architectures at scale. In standard transformer models, the attention mechanism computes correlations between a query vector $q$ and a sequence of key-value pairs $(K, V)$ to produce an output. The computational bottleneck arises because the naive implementation requires storing all $n$ previously seen tokens (the "KV cache"), leading to linear space complexity $O(n)$ and quadratic time complexity $O(n^2)$ relative to the context length.

While recent works (e.g., Haris et al., COLT'25; Kochetkova et al., NeurIPS'25) formalized this as a streaming problem and provided algorithms based on discrepancy theory, a significant gap remained between their upper bounds (algorithmic space complexity) and lower bounds (information-theoretic limits). Specifically, previous algorithms exhibited space usage scaling linearly with the inverse precision parameter $1/\epsilon$ , whereas the lower bounds did not reflect this dependence, leaving the true complexity of the problem unresolved.

The paper focuses on two specific regimes:

High Temperature Regime (Small $r$ ): Where the radius $r$ of key/query vectors is small ( $r = (\log n)^{o(1)}$ ).
Low Temperature Regime (Large $r$ ): Where $r^2 = \Omega(\log n)$ , leading to exponential dependence on $r$ .

The goal is to determine the tight space complexity for approximating the normalizing constant (Kernel Density Estimation, KDE) and the full attention vector in a streaming setting, where data arrives sequentially and must be processed with limited memory.

Methodology

The authors employ a dual approach, developing novel algorithmic upper bounds and matching information-theoretic lower bounds.

1. Algorithmic Upper Bounds

The proposed algorithms achieve tight bounds by integrating three distinct methodological pillars:

Discrepancy-Based Coreset Constructions: Utilizing results from Charikar-Kapralov-Waingarten (CKW'24) and Banaszczyk's theorem, the authors construct small subsets (coresets) of the data that approximate the kernel density estimation (KDE) value for any query. This is particularly effective for the "tail" of the exponential kernel expansion.
The Polynomial Method: Drawing on the Fast Multipole Method (Greengard-Rokhlin '87) and recent advances (Alman-Song '23), the authors use truncated polynomial expansions (Taylor or Hermite) to approximate the "head" of the kernel. They maintain a sketch of low-degree moments of the dataset to approximate the low-degree part of the exponential function.
Space Partitioning and Recentering: For the low-temperature regime (large $r$ $r$ ), the authors introduce a recursive Partition procedure. This decomposes the dataset into:
- Pseudo-random components: Where the data is spread out, allowing standard discrepancy minimization to work efficiently.
- Low-radius components: Dense clusters where the authors apply a recentering and rescaling technique. By shifting the cluster to the origin and rescaling, they reduce the effective radius of the data, significantly lowering the error term in the coreset construction compared to applying the algorithm globally.

The final streaming algorithm combines these components using a Merge-and-Reduce framework, allowing offline compression primitives to be adapted to the streaming model.

2. Lower Bound Techniques

To prove tight lower bounds, the authors reduce the INDEX problem (a communication complexity problem where Alice holds a string $x$ and Bob holds an index $i$ ) to the streaming attention problem.

Side Information Strategy: A key technical contribution is the use of "side information" in the reduction. Unlike previous reductions that failed due to high noise in the kernel values, the authors have Alice send a sketch of the low-degree part of the dataset (mimicking the upper bound's polynomial sketch).
Hermite Expansion: Instead of standard Taylor expansions, the lower bound proof utilizes Hermite polynomials. This choice is critical because it allows for a better signal-to-noise ratio when distinguishing between the presence and absence of a specific token in the dataset. The noise term (contributed by other tokens) is shown to be smaller than the signal term only when using this specific basis, enabling a tight dependence on the precision parameter $\epsilon$ and the radius $r$ .

Key Results

The paper establishes nearly tight bounds for the space complexity of streaming attention and softmax KDE, resolving the gap left by prior work.

1. High Temperature Regime (Small $r$ )

Assuming $r = (\log n)^{o(1)}$ and dimension $d = \log^a n$ (for $a > 1$ ):

Upper Bound: The problem can be solved using $\tilde{O}\left((1/\epsilon)^{1 - 1/a + o(1)}\right)$ space.
Lower Bound: Any algorithm requires $\tilde{\Omega}\left(\min\left(n, (1/\epsilon)^{1 - 1/a - o(1)}\right)\right)$ space.
Significance: This result demonstrates that the space complexity is sub-linear in $1/\epsilon$ (specifically $(1/\epsilon)^{1-1/a}$ ), improving upon the linear $O(1/\epsilon)$ dependence of previous algorithms. For example, if $d = \log^2 n$ ( $a=2$ ), the space complexity is roughly $\tilde{O}(1/\sqrt{\epsilon})$ .

2. Low Temperature Regime (Large $r$ )

Assuming $r^2 = \Omega(\log n)$ and $d = (\log n)^{1+\Omega(1)}$ :

Upper Bound: The problem can be solved using $\tilde{O}\left(\frac{e^{r^2(1+o(1))}}{\epsilon}\right)$ space.
Lower Bound: The problem requires $\tilde{\Omega}\left(\min\left(n, \frac{e^{r^2(1-o(1))}}{\epsilon}\right)\right)$ space.
Significance: This provides a nearly tight characterization of the exponential dependence on $r$ , showing that the previous bounds were not optimal regarding the exponent's constant factors.

Significance and Claims

The paper claims to provide the first nearly tight bounds for the space complexity of streaming attention across a wide range of parameters.

Closing the Gap: The work resolves the significant discrepancy between upper and lower bounds identified in prior literature (Haris et al., Kochetkova et al.), showing that the space complexity is not merely linear in $1/\epsilon$ but can be significantly sub-linear depending on the dimensionality.
Algorithmic Innovation: The paper highlights the "surprisingly tight interplay" between discrepancy theory, polynomial methods, and geometric space partitioning. The introduction of the recentering and rescaling technique for low-radius clusters is presented as a key mechanism for improving the dependence on the radius $r$ .
New Lower Bound Technique: The authors claim their technique for using the INDEX problem with large amounts of side information (specifically via Hermite expansions) is a novel contribution that may be useful for other high-dimensional geometric estimation problems.
Practical Implications: While the paper focuses on theoretical bounds, it contextualizes these results within the practical challenge of KV cache compression in Large Language Models (LLMs). By establishing that sub-linear space is theoretically achievable and quantifying the exact trade-offs with precision and dimension, the work provides a theoretical foundation for developing more efficient memory management strategies in transformers.

The authors remain modest about the scope, noting that their lower bounds for the attention problem specifically hold under the assumption that $r^2 < (\log n)/2$ in the low-temperature regime, and that extending these results to other parameter ranges remains an open problem. They also identify adaptive streaming (where inputs are generated via feedback loops) and Gaussian KDE as interesting directions for future research.