One Pass Streaming Algorithm for Super Long Token Attention Approximation in Sublinear Space

## Abstract

Attention computation takes both the time complexity of O ​ ( n 2 ) 𝑂 superscript 𝑛 2 O(n^{2}) and the space complexity of O ​ ( n 2 ) 𝑂 superscript 𝑛 2 O(n^{2}) simultaneously, which makes deploying Large Language Models (LLMs) in streaming applications that involve long contexts requiring substantial computational resources. In recent OpenAI DevDay (Nov 6, 2023), OpenAI released a new model that is able to support a 128K-long document, in our paper, we focus on the memory-efficient issue when context length n 𝑛 n is much greater than 128K ( n ≫ 2 d much-greater-than 𝑛 superscript 2 𝑑 n\gg 2^{d} ). Considering a single-layer self-attention with Query, Key, and Value matrices Q , K , V ∈ ℝ n × d 𝑄 𝐾 𝑉 superscript ℝ 𝑛 𝑑 Q,K,V\in\mathbb{R}^{n\times d} , the polynomial method approximates the attention output T ∈ ℝ n × d 𝑇 superscript ℝ 𝑛 𝑑 T\in\mathbb{R}^{n\times d} . It accomplishes this by constructing U 1 , U 2 ∈ ℝ n × t subscript 𝑈 1 subscript 𝑈 2 superscript ℝ 𝑛 𝑡 U_{1},U_{2}\in\mathbb{R}^{n\times t} to expedite attention 𝖠𝗍𝗍𝗇 ​ ( Q , K , V ) 𝖠𝗍𝗍𝗇 𝑄 𝐾 𝑉 {\sf Attn}(Q,K,V) computation within n 1 + o ​ ( 1 ) superscript 𝑛 1 𝑜 1 n^{1+o(1)} time executions. Despite this, computing the approximated attention matrix U 1 ​ U 2 ⊤ ∈ ℝ n × n subscript 𝑈 1 superscript subscript 𝑈 2 top superscript ℝ 𝑛 𝑛 U_{1}U_{2}^{\top}\in\mathbb{R}^{n\times n} still necessitates O ​ ( n 2 ) 𝑂 superscript 𝑛 2 O(n^{2}) space, leading to significant memory usage. In response to these challenges, we introduce a new algorithm that only reads one pass of the data in a streaming fashion. This method employs sublinear space o ​ ( n ) 𝑜 𝑛 o(n) to store three sketch matrices, alleviating the need for exact K , V 𝐾 𝑉 K,V storage. Notably, our algorithm exhibits exceptional memory-efficient performance with super-long tokens. As the token length n 𝑛 n increases, our error guarantee diminishes while the memory usage remains nearly constant. This unique attribute underscores the potential of our technique in efficiently handling LLMs in streaming applications.

## 1 Introduction

Large Language Models (LLMs) such as ChatGPT [], InstructGPT [], Palm [, ], BARD [], GPT-4 [], LLAMA [], LLAMA 2 [], Adobe firefly [], have revolutionized various aspects of human work. These models have shown remarkable capabilities in dialog systems [, , ], document summarization [, , , ], code completion [, , ], and question-answering [, , ]. However, their performance in these applications is often constrained by the context length.

To prepare for the coming of artificial general intelligence (AGI) [], one of the crucial bottlenecks for nowadays LLM is about how to handle super long context. In recent OpenAI DevDay (Nov 6, 2023) [] (^1^11OpenAI DevDay, Opening Keynote. [https://www.youtube.com/watch?v=U9mJuUkhUzk](https://www.youtube.com/watch?v=U9mJuUkhUzk)), OpenAI released a new model that is able to support a 128K-long document. In other words, you can feed a 300-page textbook into LLM. This is already quite surprising. However, to finally achieve AGI, we might need to feed some data that is significantly larger than the memory in a model. For example, what if we can’t even store the entire $x$ pages of a book in memory when $x$ is super large?

A longer context length allows the LLM to incorporate more information, potentially leading to more accurate and contextually appropriate responses. This increased capacity for information processing can enhance the LLM’s understanding, coherence, and contextual reasoning abilities. Therefore, to optimally utilize pretrained LLMs, it’s crucial to efficiently and accurately generate long sequences.

Despite the advantages of a long context length, LLMs, especially those based on transformers, face significant computational challenges. Inference with a long context in LLMs is computationally intensive, requiring both $O(n^{2})$ space complexity and $O(n^{2})$ time complexity to compute the attention output. This computational demand can limit the practical application of LLMs in real-world scenarios, making it a crucial area for further research and optimization.

Figure: Figure 1: Comparison between our method and previous works. On the left: vanilla attention computation []; Middle: fast attention by polynomial method []; On the right: one pass algorithm (ours).
Refer to caption: /html/2311.14652/assets/x1.png

Previous work [, ] has conducted an in-depth study on the fast approximation of attention computation within $n^{1+o(1)}$ time executions without space requirements. Below is a formal definition:

###### Definition 1.1 (Static Attention Approximation without Space Requirement [ 15 ] ) .

Let $\epsilon\in(0,1)$ denote an accuracy parameter.
Given three matrices $Q,K,V\in\mathbb{R}^{n\times d}$, the goal is to construct $T\in\mathbb{R}^{n\times d}$ such that

$$ $\displaystyle\|T-\mathsf{Attn}(Q,K,V)\|_{\infty}\leq\epsilon$ $$

where

- •
$\mathsf{Att}(Q,K,V):=D^{-1}AV$
- •
$A\in\mathbb{R}^{n\times n}$ is a square matrix $A:=\exp(QK^{\top}/d)$, here we apply $\exp()$ function entry-wisely.
- •
$D\in\mathbb{R}^{n\times n}$ is a diagonal matrix $D:=\operatorname{diag}(A{\bf 1}_{n})$ where ${\bf 1}_{n}\in\mathbb{R}^{n}$ is a length-$n$ vector where all the entries are ones.

However, the memory requirement of caching attention matrix $D^{-1}A$ for LLM’s inference is still a considerable issue that consumes $O(n^{2})$ space complexity. In this paper, we study the computation-efficiency problem in the context of transformer-based LLMs with super long context. We consider the following problem:

How can we compute the attention with super-long context in space complexity of $o(n)$?

This question is crucial as it directly relates to the computational efficiency of LLMs, particularly when dealing with super-long context lengths. In response to this question, our goal is to develop an effective streaming algorithm. We aim to define and solve the streaming version of approximate attention computation, which is a critical aspect of our research. By addressing this problem, we hope to significantly enhance the computational efficiency of transformer-based LLMs, thereby expanding their applicability in various real-world scenarios. We define the streaming version of approximate attention computation, which is also the problem we aim to solve in this paper:

###### Definition 1.2 (Streaming Attention Approximation with Sublinear in $n$ Space) .

Given $Q,K,V\in\mathbb{R}^{n\times d}$, we’re only allowed to use $o(n)$ spaces and read $Q,K,V$ in one pass, and then outputs $T\in\mathbb{R}^{n\times d}$ such that $T$ is close to $D^{-1}AV$.

###### Definition 1.1 (Static Attention Approximation without Space Requirement [ 15 ] ) .

###### Definition 1.2 (Streaming Attention Approximation with Sublinear in n 𝑛 n Space) .

### 1.1 Our Result

In our research, we tackle the challenge of efficiently calculating attention for extremely long input sequences (super-long context) with limited memory resources. Our goal is to process Query $Q$, Key $K$, and Value $V$ matrices of size $n\times d$ in a single streaming pass while utilizing only $o(n)$ space.

To address this, we propose a novel one-pass streaming algorithm (Algorithm [1](#alg1)). For $d=O(\log n)$, we first compute low-rank approximation matrices $U_{1},U_{2}\in\mathbb{R}^{n\times t}$ as in prior work [] such that $D^{-1}U_{1}U_{2}^{\top}\approx{\sf Attn}(Q,K,V)$.

Next, we introduce sketching matrices $\Phi\in\mathbb{R}^{m_{1}\times n}$ [], $\Psi\in\mathbb{R}^{m_{2}\times n}$ [] to sample $U_{1},U_{2}$ respectively, where $m_{1}=O(\epsilon_{1}^{-1}k\log n),m_{2}=\mathcal{O}(\epsilon_{2}^{-2}\log n)$ and $k$ controls sparsity.

We present our main result as follows:

###### Theorem 1.3 (Main Result, informal version of Theorem 4.1 ) .

There is a one pass streaming algorithm (Algorithm [1](#alg1)) that reads $Q,K,V\in\mathbb{R}^{n\times d}$ with $d=O(\log n)$, uses $O(\epsilon_{1}^{-1}kn^{o(1)}+\epsilon_{2}^{-2}n^{o(1)})$ spaces and outputs a matrix $T\in\mathbb{R}^{n\times d}$ such that

- •
For each $i\in[d]$, $T_{*,i}\in\mathbb{R}^{n}$ is $O(k)$-sparse column vector
- •
For each $i\in[d]$, $\|T_{i}-y_{i}\|_{2}\leq(1+\epsilon)\cdot\min_{k-\mathrm{sparse}~{}y^{\prime}}\|y^{\prime}-y_{i}\|_{2}+\epsilon_{2}$ where $y_{i}=D^{-1}AV_{*,i}$
- •
The succeed probability $0.99$.

The purpose of our work is to address the memory constraints associated with computing attention over very long sequences where the context length $n\gg 2^{d}$ (potentially infinitely long), then furthermore contribute towards more efficient and scalable transformer models, which could assist in advancing capabilities towards artificial general intelligence (AGI) []. Section [2](#S2) discusses related work that focuses on approximating attention computation. This includes prior studies on fast approximations without space requirements, which lay the groundwork for our streaming formulation. In Section [3](#S3), we outline the preliminary concepts and definitions used in our analysis. This includes problem definition, attention computation, and sketching techniques. Our key technical contributions are presented in Section [4](#S4). Here, we introduce a novel one-pass streaming algorithm for attention approximation with sublinear $o(n)$ space complexity. We also state our main theorem, which establishes performance guarantees for our proposed algorithm. Section [4](#S4) further provides a detailed proof of the main theorem. This validates that our algorithm is able to process queries, keys and values in a single streaming pass while meeting the stated approximation bounds using limited memory.

###### Theorem 1.3 (Main Result, informal version of Theorem 4.1 ) .

## 2 Related Work

In this section, we briefly review three topics that have close connections to this paper, which are Attention Theory, Streaming Algorithm and Improving LLM’s Utilization of Long Text.

#### Attention Theory

Numerous recent studies have explored attention computation in Large Language Models (LLMs) [, , , , , , , , , , , , , , , , , , , , , , , , , , , ]. Some have focused on the benefits of multiple attention heads, showing improved optimization and generalization []. Others have proposed methods like Deja Vu to reduce computational cost during inference without sacrificing quality or learning ability []. Formal analyses have examined lower and upper bounds for attention computation [, , ], while dynamic attention computation has also been investigated []. Regression problems within in-context learning for LLMs have been addressed, with a unique approach using matrix formulation []. These studies collectively contribute to our understanding of attention models and their optimization, generalization, and efficiency.

#### Streaming Algorithm

Streaming algorithms have been extensively studied in graph problems [, , , , , , , , , , , , , , , , , , ], spanning trees [], convex programming [, ], cardinality estimation [], frequency estimation [, ], sampler data structures [], heavy hitter detection [], and sparse recovery []. These studies focus on developing efficient algorithms for various problem domains, such as processing massive graphs, constructing spanning trees, optimizing convex programs, estimating cardinality and frequency, designing sampler data structures, detecting heavy hitters, and recovering sparse signals. The advancements in these areas contribute to the development of efficient and scalable algorithms for real-time analysis of streaming data.

#### Improving LLMs’ Utilization of Long Text

Extensive research has been conducted on the application of Large Language Models (LLMs) to lengthy texts [, , , , , , , , , ]. These studies aim to optimize LLMs to effectively capture and utilize the content within longer contexts, rather than treating them solely as inputs. However, despite advancements in these two directions, competent utilization of lengthy contexts within LLMs remains a challenge, as highlighted by recent works [, ].

The effective usage of prolonged contexts poses a significant challenge in the development and application of LLMs. While research has focused on improving LLMs’ understanding of longer texts, successfully leveraging this understanding for improved performance is not guaranteed. The challenge lies in effectively incorporating and utilizing the information contained within lengthy contexts, ensuring that LLMs can make accurate and meaningful predictions based on this additional context.

## 3 Preliminary

#### Notations.

We use $\operatorname{poly}(n)$ to denote $O(n^{c})$ where $c\geq 1$ is some constant.

For a vector $x\in\mathbb{R}^{n}$, we use $\|x\|_{2}$ to denote its $\ell_{2}$ norm.

We use $\|A\|_{\infty}$ to denote the $\ell_{\infty}$ norm of $A$, i.e., $\|A\|_{\infty}:=\max_{i,j}|A_{i,j}|$.

We use $\|A\|$ to denote the spectral norm of a matrix. Then it is obvious that $\|A\|\geq\max_{j}\|A_{*,j}\|_{2}$.

For a vector $x\in\mathbb{R}^{n}$, we say $x$ is $k$-sparse if and only there are $k$ nonzero entries in $x$.

For a vector $w\in\mathbb{R}^{n}$, we use $\operatorname{diag}(w)\in\mathbb{R}^{n\times n}$ to denote a diagonal matrix where the $i,i$-th entry on diagonal is $w_{i}$.

We use $\Pr[]$ to denote the probability.

### 3.1 Polynomial Method

We state a tool from previous work.

###### Lemma 3.1 (Error Approximation, Lemma 3.6 in [ 15 ] ) .

if the following conditions

- •
Let $Q,K,V\in\mathbb{R}^{n\times d}$
- •
Let $d=O(\log n)$
- •
Let $B=O(\sqrt{\log n})$
- •
Let $\|Q\|_{\infty},\|K\|_{\infty},\|V\|_{\infty}\leq B$
- •
Let $A:=\exp(QK^{\top}/d)$
- •
$D:=\operatorname{diag}(A{\bf 1}_{n})$

Then, there are matrices $U_{1},U_{2}\in\mathbb{R}^{n\times t}$ such that

- •
Part 1. $t=n^{o(1)}$
- •
Part 2. For $i$-th row in $U_{1}$, we can construct it based on $i$-th row in $Q$ in $O(t+d)$ time.
- •
Part 3. For $i$-th row in $U_{2}$, we can construct it based on $i$-th row in $K$ in $O(t+d)$ time.
- •
Part 4. Let $\widetilde{A}:=U_{1}U_{2}^{\top}$, let $\widetilde{D}:=\operatorname{diag}(\widetilde{A}{\bf 1}_{n})$, then
$\displaystyle\|D^{-1}Av-\widetilde{D}^{-1}\widetilde{A}v\|_{\infty}\leq 1/\operatorname{poly}(n)$

###### Lemma 3.1 (Error Approximation, Lemma 3.6 in [ 15 ] ) .

### 3.2 Sketching Matrices

###### Definition 3.2 ( $k$ -wise independence) .

We say $\mathcal{H}=\{h:[m]\to[l]\}$ is a $k$-wise independent hash family if $\forall i_{1}\neq i_{2}\neq\cdots\neq i_{k}\in[n]$ and $\forall j_{1},\cdots,j_{k}\in[l]$,

$$ $\displaystyle\Pr_{h\in\mathcal{H}}[h(i_{1})=j_{1}\land\cdots\land h(i_{k})=j_{k}]=\frac{1}{l^{k}}.$ $$

###### Definition 3.3 (Random Gaussian matrix) .

We say $\Psi\in\mathbb{R}^{m\times n}$ is a random Gaussian matrix if all entries are sampled from $\mathcal{N}(0,1/m)$ independently.

###### Definition 3.4 (AMS sketch matrix [ 13 ] ) .

Let $h_{1},h_{2},\cdots,h_{m}$ be $m$ random hash functions picking from a 4-wise independent hash family $\mathcal{H}=\{h:[n]\to\{-\frac{1}{\sqrt{m}},+\frac{1}{\sqrt{m}}\}\}$. Then $\Psi\in\mathbb{R}^{m\times n}$ is a AMS sketch matrix if we set $\Psi_{i,j}=h_{i}(j)$.

Note that in streaming setting, we never need to explicit write the $m\times n$ matrix. That is too expensive since it takes $\Omega(n)$ space. It is well-known that in the streaming area, we only need to store those $m$ hash functions, and each hash function only needs $O(\log n)$-bits.
Thus, the overall store for storing $\Phi$ is just $O(m\log n)$ bits

###### Definition 3.2 ( k 𝑘 k -wise independence) .

###### Definition 3.3 (Random Gaussian matrix) .

###### Definition 3.4 (AMS sketch matrix [ 13 ] ) .

### 3.3 Approximate Matrix Product

###### Lemma 3.5 (Johnson–Lindenstrauss lemma, folklore, [ 50 ] ) .

Let $m_{2}=O(\epsilon^{-2}\log(1/\delta))$.
For any fixed vectors $u$ and $v\in\mathbb{R}^{n}$, let $\Psi\in\mathbb{R}^{m_{2}\times n}$ denote a random Gaussian/AMS matrix, we have

$$ $\displaystyle\Pr[|\langle\Psi u,\Psi v\rangle-\langle u,v\rangle|\leq\epsilon\|u\|_{2}\|v\|_{2}]\geq 1-\delta$ $$

###### Lemma 3.6 .

If the following conditions hold

- •
Let $\delta\in(0,1)$ denote the failure probability
- •
Let $\epsilon_{2}\in(0,1)$ denote the accuracy parameter
- •
Let $m_{2}=O(\epsilon_{2}^{-2}\log(nd/\delta))$.
- •
Let $\|V\|\leq 1/\sqrt{n}$.

Then we have: with probability $1-\delta$

- •
Part 1. for all $j\in[n],i\in[d]$
$\displaystyle|(\widetilde{D}^{-1}U_{1}U_{2}^{\top}V)_{j,i}-(\widetilde{D}^{-1}U_{1}U_{2}^{\top}\Psi^{\top}\Psi V)_{j,i}|\leq\epsilon_{2}/\sqrt{n}$
- •
Part 2. for all $i\in[d]$, we have
$\displaystyle\|\widetilde{D}^{-1}U_{1}U_{2}^{\top}V_{*,i}-\widetilde{D}^{-1}U_{1}U_{2}^{\top}\Psi^{\top}\Psi V_{*,i}\|_{2}\leq\epsilon_{2}$

###### Proof.

First of all, $\|V\|\leq 1/\sqrt{n}$ directly implies that

$$ $\displaystyle\max_{i\in[d]}\|V_{*,i}\|_{2}\leq 1/\sqrt{n}$ (1) $$

The proof follows from applying Lemma [3.5](#S3.Thmtheorem5) and applying a union bound over $nd$ coordinates.

Proof of Part 1.
For each $j\in[n]$, for each $i\in[d]$, we can show that

$$ $\displaystyle~{}|(\widetilde{D}^{-1}U_{1}U_{2}^{\top}V)_{j,i}-(\widetilde{D}^{-1}U_{1}U_{2}^{\top}\Psi^{\top}\Psi V)_{j,i}|$ $\displaystyle\leq$ $\displaystyle~{}\epsilon_{2}\cdot\|(\widetilde{D}^{-1}U_{1}U_{2}^{\top})_{j,*}\|_{2}\cdot\|V_{*,i}\|_{2}$ $\displaystyle\leq$ $\displaystyle~{}\epsilon_{2}\cdot\|V_{*,i}\|_{2}$ $\displaystyle\leq$ $\displaystyle~{}\epsilon_{2}\cdot\frac{1}{\sqrt{n}}$ $$

where the first step follows from Lemma [3.5](#S3.Thmtheorem5), the second step follows from $\|(\widetilde{D}^{-1}U_{1}U_{2}^{\top})_{j,*}\|_{2}\leq\|(\widetilde{D}^{-1}U_{1}U_{2}^{\top})_{j,*}\|_{1}=1$, the third step follows from Eq. ([1](#S3.E1)).

Proof of Part 2.
For each $i\in[d]$, we can show that

$$ $\displaystyle~{}\|\widetilde{D}^{-1}U_{1}U_{2}^{\top}V_{*,i}-\widetilde{D}^{-1}U_{1}U_{2}^{\top}\Psi^{\top}\Psi V_{*,i}\|_{2}$ $\displaystyle=$ $\displaystyle~{}\|(\widetilde{D}^{-1}U_{1}U_{2}^{\top}V)_{*,i}-(\widetilde{D}^{-1}U_{1}U_{2}^{\top}\Psi^{\top}\Psi V)_{*,i}\|_{2}$ $\displaystyle\leq$ $\displaystyle~{}(n\cdot(\epsilon_{2}/\sqrt{n})^{2})^{1/2}$ $\displaystyle\leq$ $\displaystyle~{}\epsilon_{2}$ $$

where the first step follows from $AB_{*,i}=(AB)_{*,i}$ for all $i\in[d]$, second step follows from Part 1, the third step follows from definition of $\ell_{2}$ norm.
∎

###### Lemma 3.5 (Johnson–Lindenstrauss lemma, folklore, [ 50 ] ) .

###### Lemma 3.6 .

###### Proof.

### 3.4 Sparse Recovery

We state a sparse recovery tool from previous work [].

###### Lemma 3.7 (Sparse Recovery, Theorem 1.1 in [ 66 ] ) .

For any vector $x\in\mathbb{R}^{n}$, there is an oblivious sketching matrices $\Phi\in\mathbb{R}^{m_{1}\times n}$ such that

- •
Let $k$ denote a positive integer.
- •
$m_{1}=O(\epsilon_{1}^{-1}k\log n)$
- •
The encoding/update time(or the column sparsity of $\Phi$) is $O(\log n)$
–
In particular, computing $\Phi e_{i}\Delta$ takes $O(\log n)$ for any scalar $\Delta\in\mathbb{R}$, and one-hot vector $e_{i}\in\mathbb{R}^{n}$.
–
For convenient of later analysis, we use $z=\Phi x$.
–
The space is to store $\Phi$ is $O(m)$ bits
- •
The decoding/recover time is $O(m_{1}\log n)$
- •
The algorithm is able to output a $k$-sparse vector $x^{\prime}\in\mathbb{R}^{n}$ such that
$\displaystyle\|x^{\prime}-x\|_{2}\leq(1+\epsilon_{1})\min_{k-\mathrm{sparse}~{}x_{k}}\|x_{k}-x\|_{2}$
- •
The succeed probability is $0.999$

###### Lemma 3.7 (Sparse Recovery, Theorem 1.1 in [ 66 ] ) .

## 4 Analysis

Figure: Algorithm 1 Our One-Pass Streaming Algorithm for matrices $Q,K\in\mathbb{R}^{n\times d}$ and $V\in\mathbb{R}^{n\times d}$. The goal of this algorithm is to provide a $k$-sparse approximation to column of $Y=D^{-1}\exp(QK^{\top})V\in\mathbb{R}^{n\times d}$.

We present the main result of this paper.

###### Theorem 4.1 (Main Result, formal version of Theorem 1.3 ) .

If the following conditions hold

- •
Let $d=O(\log n)$
- •
Let $B=O(\sqrt{\log n})$
- •
Let $\|Q\|_{\infty}\leq B$, $\|K\|_{\infty}\leq B$, $\|V\|\leq 1/\sqrt{n}$
- •
Let $A:=\exp(QK^{\top}/d)\in\mathbb{R}^{n\times n}$
- •
Let $D:=\operatorname{diag}(A{\bf 1}_{n})\in\mathbb{R}^{n\times n}$

There is a one pass streaming algorithm (Algorithm [1](#alg1)) that reads $Q,K,V\in\mathbb{R}^{n\times d}$ uses

$$ $\displaystyle O(\epsilon_{1}^{-1}kn^{o(1)}+\epsilon_{2}^{-2}n^{o(1)})$ $$

spaces and outputs a matrix $T\in\mathbb{R}^{n\times d}$ such that

- •
For each $i\in[d]$, $T_{*,i}\in\mathbb{R}^{n}$ is $O(k)$-sparse column vector
- •
For each $i\in[d]$, $\|T_{i}-y_{i}\|_{2}\leq(1+\epsilon_{1})\cdot\min_{k-\mathrm{sparse}~{}y^{\prime}}\|y^{\prime}-y_{i}\|_{2}+\epsilon_{2}$ where $y_{i}=D^{-1}AV_{*,i}$
- •
The succeed probability $0.99$.
- •
The decoding time is $O(\epsilon_{1}^{-1}kn^{o(1)})$.

###### Proof.

The streaming will be able to construct sketch $Z\in\mathbb{R}^{m_{1}\times d}$ which is

$$ $\displaystyle Z=$ $\displaystyle~{}\operatorname{sk}(D^{-1}U_{1})\operatorname{sk}(U_{2})^{\top}\operatorname{sk}(V)$ $\displaystyle=$ $\displaystyle~{}\Phi D^{-1}U_{1}U_{2}\Psi^{\top}\Psi V$ $$

Running Lemma [3.7](#S3.Thmtheorem7) on $Z$ is essentially, doing sparse recovery for $D^{-1}U_{1}U_{2}\Psi^{\top}\Psi V$.

Since $D^{-1}U_{1}U_{2}\Psi^{\top}\Psi V$ is close to $D^{-1}U_{1}U_{2}V$, thus, we can finally show the error guarantees for $D^{-1}U_{1}U_{2}V$.

Proof of Space Requirement.

From the algorithm it is easy to see, the space is coming from two parts

- •
$O(m_{1}t)$ spaces for object $\operatorname{sk}(D^{-1}U_{1})$
- •
$O(m_{2}t)$ spaces for object $\operatorname{sk}(U_{2})$

From Lemma [3.1](#S3.Thmtheorem1), we have

$$ $\displaystyle t=n^{o(1)}$ $$

From Lemma [3.6](#S3.Thmtheorem6), we have

$$ $\displaystyle m_{2}=O(\epsilon_{2}^{-2}\log(n))$ $$

From Lemma [3.7](#S3.Thmtheorem7), we have

$$ $\displaystyle m_{1}=O(\epsilon_{1}^{-1}k\log n)$ $$

Thus, total space is

$$ $\displaystyle O(m_{1}t+m_{2}t)=O(\epsilon_{1}^{-1}kn^{o(1)}+\epsilon_{2}^{-2}n^{o(1)}).$ $$

Proof of Decoding Time.

The decoding time is directly following from Lemma [3.7](#S3.Thmtheorem7), it is

$$ $\displaystyle O(m_{1}\log n)=O(\epsilon_{1}^{-1}kn^{o(1)}\log n)=O(\epsilon_{1}^{-1}kn^{o(1)}).$ $$

where the first step follows from choice of $m_{1}$, the last step follows from $O(\log n)=O(n^{o(1)})$.

Proof of Error Guarantees.

To finish the proofs, we define a list of variables

- •
$y_{i}=D^{-1}AV_{*,i}\in\mathbb{R}^{n}$
- •
$\widetilde{y}_{i}=\widetilde{D}^{-1}\widetilde{A}V_{*,i}\in\mathbb{R}^{n}$
- •
$\widehat{y}_{i}=\widetilde{D}^{-1}\widetilde{A}\Psi^{\top}\Psi V_{*,i}\in\mathbb{R}^{n}$
- •
Let $\xi_{1}$ be the value that $\|y_{i}-\widetilde{y}_{i}\|_{2}\leq\xi_{1}$ ($\xi_{1}=1/\operatorname{poly}(n)$, by Part 4 of Lemma [3.1](#S3.Thmtheorem1))
- •
Let $\xi_{1}$ be the value that $\|\widetilde{y}_{i}-\widehat{y}_{i}\|_{2}\leq\xi_{2}$ ($\xi_{2}=\epsilon_{2}$, by Part 2 of Lemma [3.6](#S3.Thmtheorem6))

Firstly, we can show that

$$ $\displaystyle\|y_{i}-\widehat{y}_{i}\|_{2}\leq$ $\displaystyle~{}\|y_{i}-\widetilde{y}_{i}\|_{2}+\|\widetilde{y}_{i}-\widehat{y}\|_{2}$ $\displaystyle\leq$ $\displaystyle~{}\xi_{1}+\xi_{2}$ (2) $$

We can show

$$ $\displaystyle\|T_{i}-y_{i}\|_{2}\leq$ $\displaystyle~{}\|T_{i}-\widehat{y}_{i}\|_{2}+\|\widehat{y}_{i}-y_{i}\|_{2}$ $\displaystyle\leq$ $\displaystyle~{}\|T_{i}-\widehat{y}_{i}\|_{2}+\xi_{1}+\xi_{2}$ $\displaystyle\leq$ $\displaystyle~{}(1+\epsilon_{1})\min_{k-\mathrm{sparse}~{}y^{\prime}}\|y^{\prime}-\widehat{y}_{i}\|_{2}+\xi_{1}+\xi_{2}$ $\displaystyle\leq$ $\displaystyle~{}(1+\epsilon_{1})\min_{k-\mathrm{sparse}~{}y^{\prime}}\|y^{\prime}-y_{i}\|_{2}$ $\displaystyle~{}+(1+\epsilon)(\xi_{1}+\xi_{2})+(\xi_{1}+\xi_{2})$ $\displaystyle\leq$ $\displaystyle~{}(1+\epsilon_{1})\min_{k-\mathrm{sparse}~{}y^{\prime}}\|y^{\prime}-y_{i}\|_{2}+3(\xi_{1}+\xi_{2})$ $\displaystyle\leq$ $\displaystyle~{}(1+\epsilon_{1})\min_{k-\mathrm{sparse}~{}y^{\prime}}\|y^{\prime}-y_{i}\|_{2}+O(\epsilon_{2})$ $$

where the first step follows from triangle inequality, the second step follows from Eq. ([4](#S4.Ex24)), the third step follows from Lemma [3.7](#S3.Thmtheorem7), the fourth step follows from triangle inequality, the fifth step follows from $\epsilon_{1}\in(0,1)$ and the last step follows from $\xi_{1}=1/\operatorname{poly}(n)<\xi_{2}=\epsilon_{2}$.

Proof of Failure Probability.

The failure probability of Lemma [3.6](#S3.Thmtheorem6) is $\delta=1/\operatorname{poly}(n)$.

The failure probability of Lemma [3.7](#S3.Thmtheorem7) is 0.001.

Taking a union bound over those Lemmas, we get failure probability $0.01$ here.

In particular, the failure probability is at most

$$ $\displaystyle 0.001+1/\operatorname{poly}(n)\leq$ $\displaystyle~{}0.001+0.001$ $\displaystyle\leq$ $\displaystyle~{}0.01.$ $$

Thus, we complete the proof.
∎

###### Theorem 4.1 (Main Result, formal version of Theorem 1.3 ) .

###### Proof.

### 4.1 A General Result

We state a result for solving cross attention ($X_{1}\neq X_{2}$). Using our framework to solve self attention ($X_{1}=X_{2}$), then the algorithm will need two passes, instead of one pass.

###### Corollary 4.2 (An application of Theorem 4.1 ) .

If the following conditions hold

- •
Let $d=O(\log n)$
- •
Let $B=O(\sqrt{\log n})$
- •
Let $W_{Q},W_{K},W_{V}\in\mathbb{R}^{d\times d}$
- •
Let $X_{1},X_{2}\in\mathbb{R}^{n\times d}$
- •
Let $Q=X_{1}W_{Q}\in\mathbb{R}^{n\times d}$
- •
Let $K=X_{2}W_{K}\in\mathbb{R}^{n\times d}$
- •
Let $V=X_{2}W_{V}\in\mathbb{R}^{n\times d}$
- •
Let $\|Q\|_{\infty}\leq B$, $\|K\|_{\infty}\leq B$, $\|V\|\leq 1/\sqrt{n}$
- •
Let $A:=\exp(QK^{\top}/d)\in\mathbb{R}^{n\times n}$
- •
Let $D:=\operatorname{diag}(A{\bf 1}_{n})\in\mathbb{R}^{n\times n}$

There is a one pass streaming algorithm (Algorithm [2](#alg2)) that reads $X\in\mathbb{R}^{n\times d}$, $W_{Q},W_{K},W_{V}\in\mathbb{R}^{n\times d}$ uses

$$ $\displaystyle O(\epsilon_{1}^{-1}kn^{o(1)}+\epsilon_{2}^{-2}n^{o(1)})$ $$

spaces and outputs a matrix $T\in\mathbb{R}^{n\times d}$ such that

- •
For each $i\in[d]$, $T_{*,i}\in\mathbb{R}^{n}$ is $O(k)$-sparse column vector
- •
For each $i\in[d]$, $\|T_{i}-y_{i}\|_{2}\leq(1+\epsilon_{1})\cdot\min_{k-\mathrm{sparse}~{}y^{\prime}}\|y^{\prime}-y_{i}\|_{2}+\epsilon_{2}$ where $y_{i}=D^{-1}AV_{*,i}$
- •
The succeed probability $0.99$.
- •
The decoding time is $O(\epsilon_{1}^{-1}kn^{o(1)})$.

###### Proof.

The proofs are similar to Theorem [4.1](#S4.Thmtheorem1). The only difference between streaming algorithms (Algorithm [1](#alg1) and Algorithm [2](#alg2)) is that, in Algorithm [2](#alg2) we don’t receive each row of $Q$ (similarly as $K,V$) on the fly anymore. Instead, we store weight $W_{Q}$, and we receive each row of $X_{1}$ on the fly. Whenever we see a row of $X_{1}$, we will compute matrix vector multiplication for that row and weight $W_{Q}$.

Similarly, we applied the same strategy for $K$ and $V$.
∎

Figure: Algorithm 2 Our Streaming Algorithm for matrices $X_{1}\in\mathbb{R}^{n\times d}$, $X_{2}\in\mathbb{R}^{n\times d},W_{Q},W_{K}\in\mathbb{R}^{d\times d}$ and $W_{V}\in\mathbb{R}^{d\times d}$. The goal of this algorithm is to provide a $k$-sparse approximation to column of $Y=D^{-1}\exp(QK^{\top})V\in\mathbb{R}^{n\times d}$.

###### Corollary 4.2 (An application of Theorem 4.1 ) .

###### Proof.

## 5 Conclusion

In this paper, our research provided a brand new theoretical perspective to solve the computational efficiency problem of the attention scheme with the super-long context, which helps to employ LLMs in streaming applications that involve long contexts. We have proposed a one-pass streaming algorithm that processes Query $Q$, Key $K$, and Value $V$ matrices of size $n\times d$ in a single streaming pass while utilizing only $o(n)$ space. Our approach, which introduces the sketching method and the sparse recovery method to approximate the attention computation, plays a vital role in enabling the deployment of LLMs in various applications and contributes to the wider objective of advancing capabilities towards Artificial General Intelligence (AGI).

## Peer Reviews

confidence: 3
contribution: 1
initial_rating: 3
presentation: 3
questions: N/A
soundness: 3
strengths: The authors address a fundamental problem that is important. They produce an algorithm with clear guarantees. Furthermore most of the assumptions are clearly stated, the proofs easy to follow and the pseudocode is clean and easy to understand.
summary: The paper proposes a new framework for doing one pass streaming algorithm instead of the quadratic attention. This is useful when the sequence length $n$ is much larger than the dimension $d$ (i.e. $n\geq 2^d$).  The method relies heavily on the polynomial method proposed by Alman and Song which shows how we can compute matrices $U_1,U_2 $ such that $D^{-1} U_1^TU_2$ is approximately equal to the full attention function. They effectively employ a strategy whereby they sketch matrix U and then using sparse recovery tools to recover the original vector. Since this can be done in a  streaming fashion, one can get a substantially smaller algorithm with very few caveats.
weaknesses: The author doesn't introduce any new techniques but this seems to be the application of the sketching machinery on top of the algorithm of Alman and Song.  

The resulting algorithm has small asymptotic complexity but seems to be in extremely impractical in several ways. First it relies on $n$ being large $n \geq 2^d $ where as in practice $d \sim 10^4$ for even small transformers. As a result this assumption is quite strong. Secondly , the type of operations necessary are not easy to implement fast in modern hardware.

---

confidence: 5
contribution: 1
initial_rating: 1
presentation: 2
questions: --
soundness: 2
strengths: --
summary: In this work, the authors study improving the time and space complexity of the attention mechanism in the non-causal setting. They give a single pass streaming algorithm that given the query, key and value vectors approximates the attention output using tools from randomized linear algebra and compressed sensing. At a high level, the algorithm computes sketches of low rank matrices $U_1, U_2 \in \mathbb{R}^{n \times k}$ so that sparse approximations for the columns of the attention output $\exp(QK^T / d)V$ can be computed using the sketches of the low rank matrices. The space required by the algorithm is $n^{o(1)}$ and the time to process the streams is $n^{1+o(1)}$.
weaknesses: 1. Attention computation can be done without using $O(n^2)$ space. See Rabe and Staats "Self-Attention does not need $O(n^2)$ memory".
2. The version of the problem being studied i.e., without using causal masking is not super relevant to the current LLMs which all use causal masking.
3. As stated, the results only output at most $k \cdot d$ nonzero entries for the entire matrix $V$ which has $n$ rows. Since we expect each of the rows in the full attention output to have similar norms, many useful rows are completely marked to zero by this algorithm unless $k \ge n/d$ at which point the sketches stored by the algorithm are linear in size, at which point the claim of streaming algorithm does not make sense and the algorithm is essentially no different from Alman and Song.
4. No experimental verification of the ideas. At least showing how the algorithm works on small instances, even without implementing the additional sparse recovery steps would have been useful.

---

confidence: 3
contribution: 1
initial_rating: 3
presentation: 1
questions: See weaknesses.
soundness: 1
strengths: Addressing LLMs with very large context length appears important and timely.
summary: The paper presents a one pass streaming algorithm for computing attention in sublinear space when d is logarithmic.
weaknesses: - The writing is poor, the paper is littered with typos and the general structure is very hard to follow. I found the paper very hard to read. Other comments aside, it cannot be published in its current state.
- I find it hard to evaluate the novelty of this paper. The related work section seems to be just a lump of citations, and it is not clear what is the actual contribution. It looks like the authors simply joined existing techniques to get the results.
- I am not sure about the motivation for a streaming computation. In what scenario will the computation be executed in a streaming fashion (i.e., the matrices need to be stored somewhere, right?)
- There is no experimental evaluation.

Meta-review: The paper proposes a streaming algorithm for Attention computation. As mentioned by the reviewers the novelty is low and the paper cannot be accepted at this point.

Decision: Reject