One Pass Streaming Algorithm for Super Long Token Attention
Approximation in Sublinear Space
- URL: http://arxiv.org/abs/2311.14652v2
- Date: Mon, 5 Feb 2024 18:30:30 GMT
- Title: One Pass Streaming Algorithm for Super Long Token Attention
Approximation in Sublinear Space
- Authors: Raghav Addanki, Chenyang Li, Zhao Song, Chiwun Yang
- Abstract summary: Attention computation takes both the time complexity of $O(n2)$ and the space complexity of $O(n2)$ simultaneously.
We introduce a new algorithm that only reads one pass of data in a streaming fashion.
Notably, our algorithm exhibits exceptional memory-efficient performance with super-long tokens.
- Score: 11.735802740426294
- License: http://creativecommons.org/licenses/by-nc-sa/4.0/
- Abstract: Attention computation takes both the time complexity of $O(n^2)$ and the
space complexity of $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$ is much greater than 128K ($n \gg 2^d$). Considering a single-layer
self-attention with Query, Key, and Value matrices $Q, K, V \in \mathbb{R}^{n
\times d}$, the polynomial method approximates the attention output $T \in
\mathbb{R}^{n \times d}$. It accomplishes this by constructing $U_1, U_2 \in
\mathbb{R}^{n \times t}$ to expedite attention ${\sf Attn}(Q, K, V)$
computation within $n^{1+o(1)}$ time executions. Despite this, computing the
approximated attention matrix $U_1U_2^\top \in \mathbb{R}^{n \times n}$ still
necessitates $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)$ to
store three sketch matrices, alleviating the need for exact $K, V$ storage.
Notably, our algorithm exhibits exceptional memory-efficient performance with
super-long tokens. As the token length $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.
Related papers
- Provably Efficient Reinforcement Learning with Multinomial Logit Function Approximation [67.8414514524356]
We study a new class of MDPs that employs multinomial logit (MNL) function approximation to ensure valid probability distributions over the state space.
Despite its benefits, introducing non-linear function approximation raises significant challenges in both computational and statistical efficiency.
arXiv Detail & Related papers (2024-05-27T11:31:54Z) - Conv-Basis: A New Paradigm for Efficient Attention Inference and Gradient Computation in Transformers [27.54512534985192]
We leverage the convolution-like structure of attention matrices to develop an efficient approximation method for attention using convolution matrices.
Our algorithm achieves almost linear time, i.e., $n1+o(1)$.
This work provides a new paradigm for accelerating attention computation in transformers to enable their application to longer contexts.
arXiv Detail & Related papers (2024-05-08T17:11:38Z) - Fast $(1+\varepsilon)$-Approximation Algorithms for Binary Matrix
Factorization [54.29685789885059]
We introduce efficient $(1+varepsilon)$-approximation algorithms for the binary matrix factorization (BMF) problem.
The goal is to approximate $mathbfA$ as a product of low-rank factors.
Our techniques generalize to other common variants of the BMF problem.
arXiv Detail & Related papers (2023-06-02T18:55:27Z) - Randomized and Deterministic Attention Sparsification Algorithms for
Over-parameterized Feature Dimension [18.57735939471469]
We consider the sparsification of the attention problem.
For any super large feature dimension, we can reduce it down to the size nearly linear in length of sentence.
arXiv Detail & Related papers (2023-04-10T05:52:38Z) - Fast Attention Requires Bounded Entries [19.17278873525312]
inner product attention computation is a fundamental task for training large language models such as Transformer, GPT-1, BERT, GPT-2, GPT-3 and ChatGPT.
We investigate whether faster algorithms are possible by implicitly making use of the matrix $A$.
This gives a theoretical explanation for the phenomenon observed in practice that attention computation is much more efficient when the input matrices have smaller entries.
arXiv Detail & Related papers (2023-02-26T02:42:39Z) - Simplifying and Understanding State Space Models with Diagonal Linear
RNNs [56.33053691749856]
This work disposes of the discretization step, and proposes a model based on vanilla Diagonal Linear RNNs.
We empirically show that, despite being conceptually much simpler, $mathrmDLR$ is as performant as previously-proposed SSMs.
We also characterize the expressivity of SSMs and attention-based models via a suite of $13$ synthetic sequence-to-sequence tasks.
arXiv Detail & Related papers (2022-12-01T18:53:06Z) - Sketching Algorithms and Lower Bounds for Ridge Regression [65.0720777731368]
We give a sketching-based iterative algorithm that computes $1+varepsilon$ approximate solutions for the ridge regression problem.
We also show that this algorithm can be used to give faster algorithms for kernel ridge regression.
arXiv Detail & Related papers (2022-04-13T22:18:47Z) - Learning a Latent Simplex in Input-Sparsity Time [58.30321592603066]
We consider the problem of learning a latent $k$-vertex simplex $KsubsetmathbbRdtimes n$, given access to $AinmathbbRdtimes n$.
We show that the dependence on $k$ in the running time is unnecessary given a natural assumption about the mass of the top $k$ singular values of $A$.
arXiv Detail & Related papers (2021-05-17T16:40:48Z) - Approximate Multiplication of Sparse Matrices with Limited Space [24.517908972536432]
We develop sparse co-occuring directions, which reduces the time complexity to $widetildeOleft((nnz(X)+nnz(Y))ell+nell2right)$ in expectation.
Theoretical analysis reveals that the approximation error of our algorithm is almost the same as that of COD.
arXiv Detail & Related papers (2020-09-08T05:39:19Z) - Streaming Complexity of SVMs [110.63976030971106]
We study the space complexity of solving the bias-regularized SVM problem in the streaming model.
We show that for both problems, for dimensions of $frac1lambdaepsilon$, one can obtain streaming algorithms with spacely smaller than $frac1lambdaepsilon$.
arXiv Detail & Related papers (2020-07-07T17:10:00Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.