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

要約

アテンション計算は$O(n^2)$の時間複雑性と$O(n^2)$の空間複雑性の両方を同時に必要とするため、長いコンテキストを含むストリーミングアプリケーションに大規模言語モデル(LLM)を展開するにはかなりの計算資源が必要となる。最近のOpenAI DevDay(2023年11月6日)において、OpenAIは128K長のドキュメントをサポートできる新しいモデルをリリースしたが、本稿では、コンテキスト長$n$が128Kよりはるかに大きい場合($n ⅳ 2^d$)のメモリ効率の問題に焦点を当てる。Query、Key、Value行列$Q, K, Vを持つ1層の自己注意を考え、多項式法により注意出力$Tを近似する。これは、注目度${sf Attn}(Q, K, V)$の計算を$n^{1+o(1)}$回の実行で済むように、$U_1, U_2 ∕in ∕mathbb{R}^{n ∕t}$を構成することで実現します。にもかかわらず、近似された注意行列$U_1U_2^top ∕in ∕mathbb{R}^{n ∕times n}$の計算には$O(n^2)$の空間が必要であり、メモリ使用量が大きい。このような課題に対して、データの1パスだけをストリーミングで読み込む新しいアルゴリズムを紹介する。この手法では、3つのスケッチ行列を格納するためにサブリニアスペース$o(n)$を用い、正確な$K, V$格納の必要性を軽減する。特筆すべきは、我々のアルゴリズムが、超長いトークンに対して卓越したメモリ効率の性能を示すことである。トークンの長さ$n$が長くなるにつれて、我々のエラー保証は減少するが、メモリ使用量はほぼ一定である。このユニークな特性は、ストリーミングアプリケーションにおいてLLMを効率的に処理するための我々の技術の可能性を強調するものである。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL